[None][fix] give the MPI worker-identity gate a budget above its measured cost

[JunyiXu-nv]

Summary

The wait_shutdown=True worker-identity gate in tensorrt_llm/llmapi/mpi_session.py bounds spawn + import tensorrt_llm + the identity barrier with a single deadline. It has to: MPIPoolExecutor is constructed lazily, so the gate's first submit() is what triggers the whole worker bootstrap.

That forces one number to cover two costs with wildly different scales. Sizing it against bootstrap (low hundreds of seconds) is the only way to stop it failing healthy-but-slow pools — and it leaves the guard unable to tell a slow bootstrap from a dead pool without burning that whole budget first. Both cases produce the same 0/N valid identities text at the same wall time.

Warm, then measure. This splits the gate into two phases:

• bootstrap — n_workers collective-free probes under the existing generous, env-tunable budget, waited on in poll slices until they have all returned. All of them, not the first: ranks finish importing seconds to tens of seconds apart, and any skew left here would be paid by the barrier out of its tight budget, turning an absolute bound into a skew bound. mpi4py raises a failed MPI_Comm_spawn inside its manager thread, so the futures simply never resolve; that thread's liveness is the one positive signal separating a dead pool from a slow one, and the wait gives up on that, not on the clock.
• barrier — the existing identity barrier under a tight, fixed 60 s deadline that is now sized against what it actually measures.

Measured

20-core x86 host, real MPI (Open MPI + mpi4py 3.1.5), a worker module whose import sleeps standing in for a cold import tensorrt_llm, driving a verbatim copy of each gate:

scenario	before	after
healthy, 30 s worker import	pass	pass (31.4 s)
healthy, ranks skewed 0–6 s	pass	pass (7.4 s)
MPI_ERR_SPAWN (unspawnable pool)	300.0 s	1.0 s
worker wedged before the barrier	300.0 s	63.4 s
...workers SIGKILLed by teardown	none	all 4

Second defect fixed

That last row is a real leak. The barrier is collective, so a stalled one returns nothing, partial_identities is empty, and _teardown_unidentified_pool's SIGKILL loop is a no-op — the workers leak until job end. The bootstrap probes are collective-free, so they come back even when the barrier does not, giving teardown real PIDs to reap. Verified above: reaped=[] before, all four worker PIDs after, on the identical wedged pool.

Precisely scoped: this covers a worker that wedges after importing. One wedged inside import tensorrt_llm never returns a probe, so it stays unidentified and still leaks — unchanged from before.

Failures raise a dedicated type, not bare RuntimeError

mpi4py.MPI.Exception and concurrent.futures.BrokenExecutor both inherit from RuntimeError, and both can come out of the gate's own submit()/result(). Re-raising on the base class to preserve an already-torn-down failure would let those two past the teardown, leaving mpi_pool dangling on a session whose __init__ never finished — __del__ then blocks on a broken pool, and its abort watchdog fails on the self.comm that the raising __init__ never assigned. The gate raises _IdentityGateFailure(RuntimeError) so the re-raise catches only its own case; callers still see a RuntimeError.

Diagnostics

Every failure of the gate now carries one greppable marker, [mpi-identity-gate], and names the phase that actually failed — tracked in a local, not hard-coded, so a broken executor surfacing during the barrier is not reported as a bootstrap failure pointing on-call at TRTLLM_MPI_IDENTITY_TIMEOUT. "Never bootstrapped" and "bootstrapped but the barrier stalled" stop being the same log line. The marker lives in the exception message deliberately: that is the channel proven to reach CI logs for this path (the message text is what shows up in Jenkins output for these failures), whereas the callers' plain print() diagnostics do not survive pytest capture.

Production path unchanged

wait_shutdown=False is the production default, used by llmapi/llm.py. It is untouched:

• _collect_worker_identities has exactly one call site, guarded by if wait_shutdown: (mpi_session.py:374-375).
• An AST diff of mpi_session.py against the base confirms the only pre-existing function bodies that changed are _collect_worker_identities and _identity_barrier_timeout (now an alias returning the same value). __init__, submit, submit_sync, shutdown, _start_mpi_pool, _wait_workers_exit, _teardown_unidentified_pool, release_exit_joins, abort are byte-identical.
• test_wait_shutdown_false_never_touches_the_identity_gate asserts the constructor submits nothing on that path.

Origin

For the record and without blame: the gate arrived in #15908, which measured worker spawn + import at "~50s" and the slowest legitimate pool build at "~117s (busy node)" in the same changeset. The deadline it shipped with was below those numbers; #16971 raised it to 300 s. This change stops the two costs sharing one number at all, so raising the bootstrap budget no longer widens the guard.

Also

The test-session prefetcher now derives its outer wait from the new identity_gate_budget() (bootstrap + barrier) rather than the bootstrap half alone, so it cannot abandon a build the library still considers healthy.

Tests

Added to tests/unittest/llmapi/test_mpi_session.py, which is registered whole-file in tests/integration/test_lists/test-db/l0_a100.yml (unittest/llmapi/test_mpi_session.py ISOLATION), so CI runs them:

• test_dead_pool_fails_without_burning_the_bootstrap_budget — Blocker: dead != slow in wall time
• test_slow_bootstrap_is_not_mistaken_for_a_dead_pool — the other half
• test_warm_up_identities_are_reaped_when_the_barrier_stalls — the leak
• test_bootstrap_waits_for_every_worker_not_just_the_first — bootstrap skew, probes resolving on a timer
• test_gate_tears_down_on_runtime_errors_from_the_pool + test_mpi_and_executor_errors_really_are_runtime_errors — both arrival routes (raised from submit(), and delivered through a future so f.result() re-raises), over BrokenExecutor, plain RuntimeError and a real mpi4py.MPI.Exception
• test_pool_can_make_progress_reads_the_real_mpi4py_internals + test_mpi4py_still_exposes_the_liveness_handle — exercises the real predicate (it is a getattr chain; a renamed mpi4py internal would otherwise degrade the fast-fail to the full deadline with every test still green). The tripwire asserts every link of executor._pool.thread, not just the thread.
• test_identity_gate_marks_every_failure_path — marker + phase on both paths
• test_barrier_phase_waits_on_the_barrier_budget_not_the_bootstrap_one — asserts the deadline actually passed to futures_wait, not just the constant
• test_wait_shutdown_false_never_touches_the_identity_gate

107 passed across test_mpi_session.py, test_session_prefetcher.py, test_session_reuse.py locally. Four pre-existing tests in test_mpi_session.py fail in this environment only because they spawn subprocesses that import a tensorrt_llm older than this checkout (build mismatch, ImportError: cannot import name '_DEFAULT_IDENTITY_TIMEOUT' / kv_cache_manager_v2); they are unrelated to this change.

Mutation-checked

The two tests guarding the most consequential behaviours were verified by mutation, not just by passing:

mutation	test that catches it
restore FIRST_COMPLETED early return	test_bootstrap_waits_for_every_worker_not_just_the_first, test_slow_bootstrap_is_not_mistaken_for_a_dead_pool
hard-code phase=bootstrap again	test_gate_tears_down_on_runtime_errors_from_the_pool

Not booked to the hang program

Deliberately not tagged [TRTLLM-13409]: this is a false-positive-failure regression, not a hang.

🤖 Generated with Claude Code