[TRTLLM-9144][fix] enhance RPC robustness

[Superjomn]

Summary by CodeRabbit

Release Notes

• Bug Fixes

  • Fix the random hang issue
    • Fix the ZeroMqQueue.get_nonblock_async, which may lose some messages
    • Move all the socket operations to a central thread to avoid the zmq thread-safe issue
  • Improved RPC startup reliability with enhanced event loop initialization and timing adjustments.
  • Enhanced shutdown behavior to properly cancel pending requests instead of waiting indefinitely.
  • Better error handling and logging across RPC communication layers.

• Tests

  • Enabled additional test coverage to ensure improved reliability of RPC functionality.

Description

Test Coverage

PR Checklist

Please review the following before submitting your PR:

• PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.

• PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.

• Test cases are provided for new code paths (see test instructions)

• Any new dependencies have been scanned for license and vulnerabilities

• CODEOWNERS updated if ownership changes

• Documentation updated as needed

• The reviewers assigned automatically/manually are appropriate for the PR.

• [x] Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

/bot [-h] ['run', 'kill', 'skip', 'reuse-pipeline'] ...

Provide a user friendly way for developers to interact with a Jenkins server.

Run /bot [-h|--help] to print this help message.

See details below for each supported subcommand.

run [--reuse-test (optional)pipeline-id --disable-fail-fast --skip-test --stage-list "A10-PyTorch-1, xxx" --gpu-type "A30, H100_PCIe" --test-backend "pytorch, cpp" --add-multi-gpu-test --only-multi-gpu-test --disable-multi-gpu-test --post-merge --extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx" --detailed-log --debug(experimental)]

Launch build/test pipelines. All previously running jobs will be killed.

--reuse-test (optional)pipeline-id (OPTIONAL) : Allow the new pipeline to reuse build artifacts and skip successful test stages from a specified pipeline or the last pipeline if no pipeline-id is indicated. If the Git commit ID has changed, this option will be always ignored. The DEFAULT behavior of the bot is to reuse build artifacts and successful test results from the last pipeline.

--disable-reuse-test (OPTIONAL) : Explicitly prevent the pipeline from reusing build artifacts and skipping successful test stages from a previous pipeline. Ensure that all builds and tests are run regardless of previous successes.

--disable-fail-fast (OPTIONAL) : Disable fail fast on build/tests/infra failures.

--skip-test (OPTIONAL) : Skip all test stages, but still run build stages, package stages and sanity check stages. Note: Does NOT update GitHub check status.

--stage-list "A10-PyTorch-1, xxx" (OPTIONAL) : Only run the specified test stages. Examples: "A10-PyTorch-1, xxx". Note: Does NOT update GitHub check status.

--gpu-type "A30, H100_PCIe" (OPTIONAL) : Only run the test stages on the specified GPU types. Examples: "A30, H100_PCIe". Note: Does NOT update GitHub check status.

--test-backend "pytorch, cpp" (OPTIONAL) : Skip test stages which don't match the specified backends. Only support [pytorch, cpp, tensorrt, triton]. Examples: "pytorch, cpp" (does not run test stages with tensorrt or triton backend). Note: Does NOT update GitHub pipeline status.

--only-multi-gpu-test (OPTIONAL) : Only run the multi-GPU tests. Note: Does NOT update GitHub check status.

--disable-multi-gpu-test (OPTIONAL) : Disable the multi-GPU tests. Note: Does NOT update GitHub check status.

--add-multi-gpu-test (OPTIONAL) : Force run the multi-GPU tests in addition to running L0 pre-merge pipeline.

--post-merge (OPTIONAL) : Run the L0 post-merge pipeline instead of the ordinary L0 pre-merge pipeline.

--extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx" (OPTIONAL) : Run the ordinary L0 pre-merge pipeline and specified test stages. Examples: --extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx".

--detailed-log (OPTIONAL) : Enable flushing out all logs to the Jenkins console. This will significantly increase the log volume and may slow down the job.

--debug (OPTIONAL) : Experimental feature. Enable access to the CI container for debugging purpose. Note: Specify exactly one stage in the stage-list parameter to access the appropriate container environment. Note: Does NOT update GitHub check status.

For guidance on mapping tests to stage names, see docs/source/reference/ci-overview.md and the scripts/test_to_stage_mapping.py helper.

kill

kill

Kill all running builds associated with pull request.

skip

skip --comment COMMENT

Skip testing for latest commit on pull request. --comment "Reason for skipping build/test" is required. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break.

reuse-pipeline

reuse-pipeline

Reuse a previous pipeline to validate current commit. This action will also kill all currently running builds associated with the pull request. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break.

[coderabbitai[bot]]

📝 Walkthrough

Walkthrough

RPC infrastructure hardening for resilience and graceful shutdown. Changes include immediate shutdown with request cancellation semantics, improved event loop startup reliability via increased delays and synchronization, centralized error handling helpers, and removal of test waivers with enhanced resource cleanup patterns.

Changes

Cohort / File(s)	Summary
RPC Client Startup Resilience tensorrt_llm/executor/rpc/rpc_client.py	Increased initial sleep delays (0.1s → 0.2s) in _ensure_event_loop and run_loop(). Added lazy response reader startup with non-blocking deferral and try/except around loop-start logic to prevent exception propagation.
RPC Server Shutdown & Cancellation tensorrt_llm/executor/rpc/rpc_server.py	Changed shutdown to immediate cancellation instead of waiting for pending work. Added internal helpers: _cancel_pending_queue_requests, _send_error_response, _handle_shutdown_request. Integrated RPCCancelled checks before and during request processing; streaming requests trigger cancellation on shutdown detection. Updated error handling to differentiate streaming vs. non-streaming responses.
RPC Proxy Synchronization & Robustness tensorrt_llm/executor/rpc_proxy.py	Added main_loop_started Event for synchronization; early response reader startup with explicit readiness check. Wrapped RPC submission with try/except for _results cleanup. Enhanced shutdown robustness via hasattr checks and exception handling around MPI session, RPC client close, and main loop cancellation with timeout-based thread join. Simplified main loop task composition.
Test Resource Cleanup & Skip Removal tests/unittest/executor/test_rpc.py, tests/unittest/executor/test_rpc_proxy.py, tests/unittest/executor/test_rpc_worker.py	Replaced context-managed client/server cleanup with explicit construction, try/finally blocks, and ordered shutdown/close sequences. Updated test assertions to expect RPCCancelled behavior on shutdown. Removed skip decorators to enable previously skipped tests; added skip marker for flaky performance test. Changed address generation to unique IPC addresses.
Test Waivers Reduction tests/integration/test_lists/waives.txt	Removed multiple SKIP entries for tests under unittest/executor and unittest/llmapi, reducing waived test exclusions.

Sequence Diagram(s)

sequenceDiagram
    participant Client
    participant Server
    participant Queue
    participant Worker

    note over Client,Worker: Normal Request Processing Flow
    Client->>Server: submit_request(req)
    Server->>Server: _handle_shutdown_request()
    alt Shutdown Active
        Server->>Queue: cancel pending requests
        Server->>Client: send RPCCancelled error
    else Normal Path
        Server->>Worker: process_request()
        Worker-->>Server: response
        Server->>Client: send_response()
    end

    note over Client,Worker: Shutdown Initiated
    Client->>Server: shutdown()
    Server->>Queue: _cancel_pending_queue_requests()
    activate Queue
    loop For each pending request
        Queue->>Client: send RPCCancelled error
    end
    deactivate Queue
    Server->>Worker: executor.shutdown(wait=False)
    Server->>Server: cancel main_loop task
    Server->>Server: stop event loop

sequenceDiagram
    participant Caller
    participant Proxy
    participant EventLoop
    participant Client

    note over Caller,Client: RPC Proxy Initialization with Sync
    Caller->>Proxy: __init__()
    Proxy->>Proxy: create main_loop_started Event
    Proxy->>Client: start_response_reader()
    activate EventLoop
    Proxy->>EventLoop: run main_loop tasks
    rect rgb(200, 220, 255)
        note over Proxy: Main loop ready
        Proxy->>Proxy: set main_loop_started
    end
    deactivate EventLoop
    Caller->>Proxy: await main_loop_started.wait(timeout)
    alt Timeout
        Proxy->>Caller: raise TimeoutError
    else Ready
        Caller->>Proxy: proceed with execute()
    end

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~60 minutes

• rpc_server.py: Requires careful analysis of new request cancellation flow, integration of _handle_shutdown_request across all request paths, and error handling differentiation between streaming/non-streaming modes
• rpc_proxy.py: Multiple synchronization mechanisms, shutdown robustness with error handling, and state management during initialization and cleanup
• rpc_client.py: Timing-sensitive logic changes and lazy initialization; verify no race conditions introduced
• Test files: Validate that new cleanup patterns properly release resources and that removed skip decorators don't mask underlying flakiness beyond the explicitly marked performance test

Pre-merge checks and finishing touches

❌ Failed checks (2 warnings)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 54.55% which is insufficient. The required threshold is 80.00%.	You can run @coderabbitai generate docstrings to improve docstring coverage.
Description check	⚠️ Warning	The PR description is entirely blank except for the template placeholders; the Description and Test Coverage sections are empty, and the checklist items are not filled out.	Provide a detailed description of the issue being fixed, explain the solution, list relevant test cases that validate the changes, and complete the PR checklist items.

✅ Passed checks (1 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title '[TRTLLM-9144][fix] enhance RPC robustness' clearly and specifically describes the main objective of the PR - enhancing robustness of the RPC system. This aligns with the comprehensive changes across rpc_client.py, rpc_server.py, and rpc_proxy.py that improve error handling, shutdown semantics, and event loop management.

✨ Finishing touches

• [ ] 📝 Generate docstrings

🧪 Generate unit tests (beta)

• [ ] Create PR with unit tests
• [ ] Post copyable unit tests in a comment

[!TIP]

📝 Customizable high-level summaries are now available in beta!

You can now customize how CodeRabbit generates the high-level summary in your pull requests — including its content, structure, tone, and formatting.

• Provide your own instructions using the high_level_summary_instructions setting.
• Format the summary however you like (bullet lists, tables, multi-section layouts, contributor stats, etc.).
• Use high_level_summary_in_walkthrough to move the summary from the description to the walkthrough section.

Example instruction:

"Divide the high-level summary into five sections:

1. 📝 Description — Summarize the main change in 50–60 words, explaining what was done.
2. 📓 References — List relevant issues, discussions, documentation, or related PRs.
3. 📦 Dependencies & Requirements — Mention any new/updated dependencies, environment variable changes, or configuration updates.
4. 📊 Contributor Summary — Include a Markdown table showing contributions: | Contributor | Lines Added | Lines Removed | Files Changed |
5. ✔️ Additional Notes — Add any extra reviewer context. Keep each section concise (under 200 words) and use bullet or numbered lists for clarity."

Note: This feature is currently in beta for Pro-tier users, and pricing will be announced later.

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

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

[Superjomn]

/bot run --stage-list "A100X-PyTorch-1"

[tensorrt-cicd]

PR_Github #23611 [ run ] triggered by Bot. Commit: 701ff1d

[tensorrt-cicd]

PR_Github #23611 [ run ] completed with state SUCCESS. Commit: 701ff1d /LLM/main/L0_MergeRequest_PR pipeline #17765 (Partly Tested) completed with status: 'FAILURE'

[Superjomn]

/bot run --only-multi-gpu-test

[tensorrt-cicd]

PR_Github #24156 [ run ] triggered by Bot. Commit: 864b3e1

[tensorrt-cicd]

PR_Github #24156 [ run ] completed with state SUCCESS. Commit: 864b3e1 /LLM/main/L0_MergeRequest_PR pipeline #18213 (Partly Tested) completed with status: 'FAILURE'

[Superjomn]

/bot run

[tensorrt-cicd]

PR_Github #24200 [ run ] triggered by Bot. Commit: 864b3e1

[tensorrt-cicd]

PR_Github #24200 [ run ] completed with state SUCCESS. Commit: 864b3e1 /LLM/main/L0_MergeRequest_PR pipeline #18248 completed with status: 'FAILURE'

[Superjomn]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #24304 [ run ] triggered by Bot. Commit: cde1bb3

[tensorrt-cicd]

PR_Github #24304 [ run ] completed with state SUCCESS. Commit: cde1bb3 /LLM/main/L0_MergeRequest_PR pipeline #18335 completed with status: 'FAILURE'

[Superjomn]

/bot run

[Superjomn]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #24899 [ run ] triggered by Bot. Commit: 9cfc32c

[tensorrt-cicd]

PR_Github #24900 [ run ] triggered by Bot. Commit: 9cfc32c

[tensorrt-cicd]

PR_Github #24899 [ run ] completed with state ABORTED. Commit: 9cfc32c

[tensorrt-cicd]

PR_Github #24900 [ run ] completed with state SUCCESS. Commit: 9cfc32c /LLM/main/L0_MergeRequest_PR pipeline #18801 completed with status: 'FAILURE'

[Superjomn]

/bot run

[tensorrt-cicd]

PR_Github #24958 [ run ] triggered by Bot. Commit: 9cfc32c

[hchings]

To verify the fix, let's see whether we're able to pass CI consistently, which will include Ray stages w/ RPC. But this might be tricky since CI itself is quite unstable lately.

I will help with:

• See whether this random issue https://nvbugspro.nvidia.com/bug/5636663 can still be repro.
• I'll also check on shutdown issues I encountered before locally.

[tensorrt-cicd]

PR_Github #24958 [ run ] completed with state SUCCESS. Commit: 9cfc32c /LLM/main/L0_MergeRequest_PR pipeline #18853 completed with status: 'FAILURE'

[Superjomn]

/bot run

[tensorrt-cicd]

PR_Github #24994 [ run ] triggered by Bot. Commit: 9cfc32c

[Superjomn]

/bot run

[tensorrt-cicd]

PR_Github #25044 [ run ] triggered by Bot. Commit: 9cfc32c

[tensorrt-cicd]

PR_Github #25044 [ run ] completed with state SUCCESS. Commit: 9cfc32c /LLM/main/L0_MergeRequest_PR pipeline #18926 completed with status: 'FAILURE'

[Superjomn]

/bot run --only-multi-gpu-test

[tensorrt-cicd]

PR_Github #25113 [ run ] triggered by Bot. Commit: 9cfc32c