[None][test] disagg startup: detect and fail fast on worker OOM/crash

[brnguyen2]

Summary

• Log work_dir and worker log paths at launch so failures are easy to locate without digging through pytest output.
• Fix worker log/config filename collision when multiple workers share port=0: use worker_index instead of port in filenames (worker_ctx_0.log, worker_ctx_1.log, etc.).
• Detect worker crashes during startup via process.poll(); persist worker logs under cwd so they survive test teardown.
• Fix workspace cleanup on overlayfs: use ignore_errors=True on shutil.rmtree since overlayfs may reject directory removal even with correct permissions.
• Fix swallowed crash exception: previously ticker.cancel() discarded any RuntimeError _tick() had raised; await ticker in a finally block so the exception propagates correctly.
• Fix watchdog exception silently leaving wait_for_disagg_server_ready polling for up to server_start_timeout: run both as asyncio tasks with asyncio.wait(FIRST_COMPLETED) so whichever finishes first (success or exception) cancels the other.
• Detect workers alive but hung (stuck in MPI cleanup after OOM) via /cluster_info: if no new workers register for 60s after the disagg server first becomes reachable, declare the cluster stuck.
• Log startup progress every 10s (elapsed, remaining, workers registered) and log total setup duration after cluster is ready.
• Lower server_start_timeout for mixed-stress from 7200s to 600s (10m) — sufficient for an 8-GPU setup with a warm Triton cache.
• Add logging to MpiSession.shutdown_abort and MpiPoolSession.shutdown to surface where MPI cleanup stalls after a worker crash.

Test plan

• Validated on B200 8-GPU machine: with GPUs 4-7 occupied by another user, the gen worker OOMs during startup. The test now fails within ~40s with a clear RuntimeError: 1 worker(s) exited during startup: worker_gen_0.log (rc=1) instead of hanging indefinitely.
• Previously the test would hang for up to 2 hours (7200s timeout) after the watchdog detected the crash, because the wait_for_disagg_server_ready coroutine was not cancelled.

Summary by CodeRabbit

• Bug Fixes

  • Improved distributed session shutdown handling with clearer progress, timeout, and failure diagnostics.
  • Enhanced startup monitoring for disaggregated deployments to detect failed workers and stalled cluster initialization earlier.
  • Prevented worker configuration and log-file conflicts when launching multiple workers.

• Tests

  • Improved cleanup of temporary test workspaces.
  • Added more reliable worker logs and startup timing information for diagnosing integration-test failures.
  • Reduced the default wait time for mixed-stress cluster startup checks.

[coderabbitai]

📝 Walkthrough

Walkthrough

MPI shutdown paths now provide lifecycle logging. Disaggregated integration workers use indexed artifact names, startup monitoring detects exited workers and registration stalls, failure logs are retained, and mixed-stress startup uses a 600-second timeout.

Changes

MPI shutdown lifecycle

Layer / File(s)	Summary
MPI shutdown logging tensorrt_llm/llmapi/mpi_session.py	Shutdown and timeout paths now log worker waits, abort attempts, clean exits, and MPI pool completion states.

Disaggregated integration startup

Layer / File(s)	Summary
Worker artifact disambiguation tests/integration/defs/disaggregated/disagg_test_utils.py	Worker indexes are threaded through context and generation worker helpers and used in configuration and log filenames.
Cluster startup watchdog and cleanup tests/integration/defs/disaggregated/test_disaggregated.py, tests/integration/defs/conftest.py	Cluster setup retains failure logs, reports worker log paths, monitors readiness and worker registration concurrently, and updates workspace cleanup.
Mixed-stress startup timing tests/integration/defs/disaggregated/test_disaggregated.py	Mixed-stress setup propagates startup settings, reports setup duration, and uses a 600-second server startup timeout.

Estimated code review effort: 3 (Moderate) | ~20 minutes

Sequence Diagram(s)

sequenceDiagram
  participant setup_disagg_cluster
  participant wait_for_disagg_server_ready
  participant worker_processes
  participant cluster_info
  setup_disagg_cluster->>worker_processes: launch context and generation workers
  setup_disagg_cluster->>wait_for_disagg_server_ready: wait for server readiness
  setup_disagg_cluster->>worker_processes: check worker process exits
  setup_disagg_cluster->>cluster_info: poll registered workers
  wait_for_disagg_server_ready-->>setup_disagg_cluster: readiness result
  cluster_info-->>setup_disagg_cluster: registration status
  setup_disagg_cluster->>worker_processes: terminate on startup failure

Loading

Suggested reviewers: xinhe-nv

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title is specific, concise, and matches the main change: fail-fast disaggregated startup on worker crashes/OOMs.
Description check	✅ Passed	The description provides a clear summary and test plan that cover the main behavior and validation for the PR.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands.}

[brnguyen2]

/bot run

[coderabbitai]

Actionable comments posted: 3

🧹 Nitpick comments (1)

tests/integration/defs/disaggregated/disagg_test_utils.py (1)

86-95: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Annotate the changed worker APIs and document the public wrapper arguments.

Add precise parameter and return types to all three signatures, plus Args/Returns sections for run_ctx_worker and run_gen_worker. As per coding guidelines, “Annotate every function” and “document public function arguments.”

Also applies to: 157-166, 184-193

🤖 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/integration/defs/disaggregated/disagg_test_utils.py` around lines 86 -
95, Annotate all three worker APIs—_run_worker, run_ctx_worker, and
run_gen_worker—with precise parameter and return types, including optional and
callback/configuration values as appropriate. Add Args and Returns sections to
the docstrings of the public wrappers run_ctx_worker and run_gen_worker,
documenting every exposed argument and the returned value. Do not alter their
runtime behavior.

Source: Coding guidelines

🤖 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`:
- Line 131: Update the shutdown logging statement in the MPI session flow to use
the f-string !r conversion flag for reason instead of calling repr(reason),
preserving the existing “Reason to shutdown” message.

In `@tests/integration/defs/disaggregated/test_disaggregated.py`:
- Around line 853-858: Update the pending-task cleanup loop to suppress only
asyncio.CancelledError when awaiting each cancelled task; remove the broad
Exception catch so any non-cancellation failure from pending tasks propagates
instead of being discarded.
- Around line 821-830: Update the worker-registration polling logic around
last_worker_count so the no-progress timeout is only evaluated while count is
below the expected worker total; once all expected workers register, stop or
bypass the timer during subsequent readiness waits. Add a mocked slow-readiness
regression in the disaggregated integration tests to verify complete
registration does not raise the no-progress RuntimeError even when final server
readiness exceeds _NO_PROGRESS_TIMEOUT.

---

Nitpick comments:
In `@tests/integration/defs/disaggregated/disagg_test_utils.py`:
- Around line 86-95: Annotate all three worker APIs—_run_worker, run_ctx_worker,
and run_gen_worker—with precise parameter and return types, including optional
and callback/configuration values as appropriate. Add Args and Returns sections
to the docstrings of the public wrappers run_ctx_worker and run_gen_worker,
documenting every exposed argument and the returned value. Do not alter their
runtime behavior.

🪄 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: 6a2fbde5-7f53-4409-9416-9a38e619b176

📥 Commits

Reviewing files that changed from the base of the PR and between 573bd5f and 51fd791.

📒 Files selected for processing (4)

• tensorrt_llm/llmapi/mpi_session.py
• tests/integration/defs/conftest.py
• tests/integration/defs/disaggregated/disagg_test_utils.py
• tests/integration/defs/disaggregated/test_disaggregated.py

[tensorrt-cicd]

PR_Github #59541 [ run ] triggered by Bot. Commit: 51fd791 Link to invocation

[brnguyen2]

/bot run

[tensorrt-cicd]

PR_Github #59555 [ run ] triggered by Bot. Commit: e38a5c5 Link to invocation

[tensorrt-cicd]

PR_Github #59541 [ run ] completed with state ABORTED. Commit: 51fd791

Link to invocation

[tensorrt-cicd]

PR_Github #59555 [ run ] completed with state SUCCESS. Commit: e38a5c5
/LLM/main/L0_MergeRequest_PR pipeline #48001 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[brnguyen2]

/bot run

[nv-xtf]

Question on the log-retention part — am I reading this right? The retained logs live under the llm_venv workspace, which the session fixture seems to rmtree at teardown, and IIUC the CI stage only uploads ${stageName}/ (the pytest --output-dir). If both hold, the full worker logs wouldn't reach the downloadable artifacts — is there an earlier collection step or another path I'm missing? If not (non-blocking, fine as a follow-up), copying them into --output-dir on failure might be all that's needed.

[BowenFu]

LGTM — disagg startup fail-fast on worker OOM via process.poll(); can't false-positive on long-lived servers, and the readiness poller is correctly cancelled/​re-raised on a crash.

[brnguyen2]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #59707 [ run ] triggered by Bot. Commit: a8b7ddd Link to invocation

[brnguyen2]

@nv-xtf Good catch. For the tests in this PR, save_log defaults to False, so worker stdout/stderr inherits fd 1/fd 2 from the parent process. In CI, the s3_output.py plugin redirects those fds via os.dup2() before each test, so subprocess output lands in the per-test stdout.log/stderr.log files under --output-dir and gets uploaded to Artifactory.

This iteration is mainly meant to make local testing easier. Explicitly copying logs into --output-dir on failure is a good follow-up once this lands.

[tensorrt-cicd]

PR_Github #59707 [ run ] completed with state SUCCESS. Commit: a8b7ddd
/LLM/main/L0_MergeRequest_PR pipeline #48134 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[brnguyen2]

/bot run

[tensorrt-cicd]

PR_Github #59765 [ run ] triggered by Bot. Commit: a8b7ddd Link to invocation

[tensorrt-cicd]

PR_Github #59765 [ run ] completed with state SUCCESS. Commit: a8b7ddd
/LLM/main/L0_MergeRequest_PR pipeline #48191 completed with status: 'SUCCESS'
Pipeline passed with automatic retried tests. Check the rerun report for details.

CI Report

Link to invocation