[TRTLLM-14040][fix] Fix multi-GPU VisualGen response IPC: REBUILD_LOCAL handle crossed a process boundary#16326
Conversation
|
/bot run --disable-fail-fast |
📝 WalkthroughWalkthroughWorkers now cache external-launch mode at startup and use it when transferring response CUDA handles. A spawned multi-worker regression test verifies HMAC-protected IPC response rebuilding with a deterministic stub pipeline. ChangesExternal Launch Handle Mode
Estimated code review effort: 3 (Moderate) | ~20 minutes 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
tests/unittest/visual_gen/test_executor_shared_tensor_ipc.py (1)
119-167: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAnnotate the new test helpers.
Add parameter and return annotations to
_StubPipelinemethods,_stub_load_pipeline,_diffusion_worker_entry, and the test method.As per coding guidelines, “Always annotate functions.”
Also applies to: 181-181
🤖 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/visual_gen/test_executor_shared_tensor_ipc.py` around lines 119 - 167, Add explicit parameter and return type annotations to all methods on _StubPipeline, plus _stub_load_pipeline, _diffusion_worker_entry, and the affected test method. Use the existing request, pipeline, output, and worker argument types where available, and annotate side-effect-only methods with None.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 `@tests/unittest/visual_gen/test_executor_shared_tensor_ipc.py`:
- Around line 181-215: Update test_spawn_worker_response_rebuilds_in_client to
create and use a local multiprocessing spawn context instead of calling
mp.set_start_method globally. Construct workers through that context, and ensure
all workers are joined in a finally block before requests_ipc and responses_ipc
are closed.
---
Nitpick comments:
In `@tests/unittest/visual_gen/test_executor_shared_tensor_ipc.py`:
- Around line 119-167: Add explicit parameter and return type annotations to all
methods on _StubPipeline, plus _stub_load_pipeline, _diffusion_worker_entry, and
the affected test method. Use the existing request, pipeline, output, and worker
argument types where available, and annotate side-effect-only methods with None.
🪄 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: e1fa0c40-c87e-468e-b538-578a14d85f77
📒 Files selected for processing (2)
tensorrt_llm/_torch/visual_gen/executor.pytests/unittest/visual_gen/test_executor_shared_tensor_ipc.py
|
PR_Github #58963 [ run ] triggered by Bot. Commit: |
18908a5 to
c0e0c33
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #58967 [ run ] triggered by Bot. Commit: |
|
PR_Github #58963 [ run ] completed with state |
…he launch site Deciding the response handle mode inside the worker via _detect_external_launch() misreads every mp-spawn worker as externally launched: run_diffusion_worker exports RANK/WORLD_SIZE for torch.distributed env:// init, after which spawn workers are indistinguishable from launcher-started ones. Rank 0 then ships a same-process (REBUILD_LOCAL) tensor handle across the process boundary and the client fails with 'REBUILD_LOCAL handle crossed a process boundary' on any single-node multi-GPU run. The launch site knows the fact the decision actually needs — whether the worker runs inside the client process — so declare it there: run_diffusion_worker and DiffusionExecutor take in_client_process (default False = cross-process handles), and only the external-launch branch that runs rank 0's worker as a client thread passes True. The worker-side environment sniffing is gone entirely; _detect_external_launch() remains only at its client-side call sites, which read the launcher-provided environment before anything mutates it. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com>
c0e0c33 to
9edf869
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #58970 [ run ] triggered by Bot. Commit: |
|
PR_Github #58967 [ run ] completed with state |
|
PR_Github #58970 [ run ] completed with state |
TestSpawnWorkerResponseHandleMode drives the real run_diffusion_worker entry,
serve loop and ZMQ response path with world_size=2 (gloo, CPU-only) and
asserts the client-side rebuild of the returned media succeeds — guarding the
in_client_process=False default for spawn workers. Only the pipeline is
stubbed, patched inside the spawned child. This is the missing combination
the REBUILD_LOCAL regression slipped through: the worker entry's
RANK/WORLD_SIZE env exports, a multi-worker world and the production handle
mode decision were each tested separately but never together. Verified
red/green: any wrongly-local handle mode fails this test with the production
error ('REBUILD_LOCAL handle crossed a process boundary').
Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com>
9edf869 to
d676a01
Compare
|
/bot kill |
|
/bot reuse-pipeline |
|
PR_Github #59087 [ reuse-pipeline ] triggered by Bot. Commit: |
|
PR_Github #59089 [ kill ] triggered by Bot. Commit: |
|
PR_Github #59087 [ reuse-pipeline ] completed with state |
|
PR_Github #59089 [ kill ] completed with state |
|
PR_Github #59090 [ run ] triggered by Bot. Commit: |
|
/bot reuse-pipeline --number 58970 |
|
PR_Github #59091 [ reuse-pipeline ] triggered by Bot. Commit: |
|
PR_Github #59090 [ run ] completed with state |
|
PR_Github #59091 [ reuse-pipeline ] completed with state |
@coderabbitai summary
Description
Since #16050, every single-node multi-GPU VisualGen run (the default
VisualGenPython API path withworld_size > 1) fails on its first generation: the client raisesREBUILD_LOCAL handle crossed a process boundary; it is only valid within the producing process.#16050 introduced a same-process (
REBUILD_LOCAL) shared-tensor handoff for external-launch mode, where rank 0's worker and the client are threads of one process. The handle mode was decided by calling_detect_external_launch()at response time — butrun_diffusion_workerhas by then already exportedRANK/WORLD_SIZEfortorch.distributedenv:// init, which makes every spawned worker indistinguishable from an externally launched one. Rank 0 therefore ships a same-process handle across the process boundary. Single-GPU runs (world_size == 1) and external-launch runs are unaffected, which is how the regression slipped through.Fix: the launch site knows the fact the decision actually needs — whether the worker runs inside the client process — so declare it there instead of sniffing the environment.
run_diffusion_workerandDiffusionExecutortakein_client_process(defaultFalse= cross-process handles); only the external-launch branch that runs rank 0's worker as a client thread passesTrue. Worker-side_detect_external_launch()usage is gone entirely; the function remains only at its client-side call sites, which read the launcher-provided environment before anything mutates it.Test Coverage
TestSpawnWorkerResponseHandleModeintests/unittest/visual_gen/test_executor_shared_tensor_ipc.py: drives the realrun_diffusion_workerentry, serve loop and ZMQ response path withworld_size=2(gloo, CPU-only; only the pipeline is stubbed, patched inside the spawned child) and asserts the client-side rebuild of the returned media succeeds — guarding thein_client_process=Falsedefault for spawn workers. This is the previously untested combination the regression needed: the worker entry's env exports, a multi-worker world and the production handle-mode decision were each covered separately but never together.cfg_size=1 x ulysses_size=8, NVFP4): generation completes and the client receives the video; before the fix the client failed on the first response.PR Checklist
PR description clearly explains what and why.
PR Follows TRT-LLM CODING GUIDELINES.
Test cases are provided for new code paths.
Any new dependencies have been scanned for license and vulnerabilities.
CODEOWNERS updated if ownership changes.
Documentation updated as needed.
Update tava architecture diagram if significant design change.
Please check this after reviewing the above items as appropriate for this PR.
GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.