Skip to content

[TRTLLM-14040][fix] Fix multi-GPU VisualGen response IPC: REBUILD_LOCAL handle crossed a process boundary#16326

Merged
luyiyun1021 merged 2 commits into
NVIDIA:mainfrom
luyiyun1021:fix/visualgen-spawn-response-handle
Jul 14, 2026
Merged

[TRTLLM-14040][fix] Fix multi-GPU VisualGen response IPC: REBUILD_LOCAL handle crossed a process boundary#16326
luyiyun1021 merged 2 commits into
NVIDIA:mainfrom
luyiyun1021:fix/visualgen-spawn-response-handle

Conversation

@luyiyun1021

@luyiyun1021 luyiyun1021 commented Jul 13, 2026

Copy link
Copy Markdown
Collaborator

@coderabbitai summary

Description

Since #16050, every single-node multi-GPU VisualGen run (the default VisualGen Python API path with world_size > 1) fails on its first generation: the client raises REBUILD_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 — but run_diffusion_worker has by then already exported RANK/WORLD_SIZE for torch.distributed env:// 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_worker and DiffusionExecutor take in_client_process (default False = cross-process handles); only the external-launch branch that runs rank 0's worker as a client thread passes True. 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

  • New TestSpawnWorkerResponseHandleMode in tests/unittest/visual_gen/test_executor_shared_tensor_ipc.py: drives the real run_diffusion_worker entry, serve loop and ZMQ response path with world_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 the in_client_process=False default 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.
  • Red/green verified: any wrongly-local handle mode fails the new test with the exact production error; with the fix, all 6 tests in the file pass.
  • End-to-end verified on 8x B200 (LTX-2 two-stage, 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.

@luyiyun1021 luyiyun1021 requested a review from a team as a code owner July 13, 2026 11:08
@luyiyun1021

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@coderabbitai

coderabbitai Bot commented Jul 13, 2026

Copy link
Copy Markdown
Contributor

Review Change Stack

📝 Walkthrough

Walkthrough

Workers 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.

Changes

External Launch Handle Mode

Layer / File(s) Summary
Cache launch mode for response handles
tensorrt_llm/_torch/visual_gen/executor.py
run_diffusion_worker detects launch mode before setting distributed environment variables, passes it to DiffusionExecutor, and uses the cached value for rank-0 response handle conversion.
Validate spawned-worker response rebuilding
tests/unittest/visual_gen/test_executor_shared_tensor_ipc.py
Adds stubbed spawned workers, HMAC-protected ZMQ queues, response tensor reconstruction assertions, and worker cleanup handling.

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

🚥 Pre-merge checks | ✅ 5
✅ Passed checks (5 passed)
Check name Status Explanation
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.
Title check ✅ Passed The title matches the main fix and follows the repository's ticket/type summary format.
Description check ✅ Passed The PR includes the required Description, Test Coverage, and PR Checklist sections with relevant details.
✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests

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

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 win

Annotate the new test helpers.

Add parameter and return annotations to _StubPipeline methods, _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

📥 Commits

Reviewing files that changed from the base of the PR and between 5603308 and 18908a5.

📒 Files selected for processing (2)
  • tensorrt_llm/_torch/visual_gen/executor.py
  • tests/unittest/visual_gen/test_executor_shared_tensor_ipc.py

Comment thread tests/unittest/visual_gen/test_executor_shared_tensor_ipc.py
@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #58963 [ run ] triggered by Bot. Commit: 18908a5 Link to invocation

@luyiyun1021 luyiyun1021 force-pushed the fix/visualgen-spawn-response-handle branch from 18908a5 to c0e0c33 Compare July 13, 2026 11:58
@luyiyun1021 luyiyun1021 changed the title [TRTLLM-14040][fix] VisualGen: mp-spawn multi-GPU response misread as same-process handoff [TRTLLM-14040][fix] VisualGen: declare same-process client at the launch site Jul 13, 2026
@luyiyun1021

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #58967 [ run ] triggered by Bot. Commit: c0e0c33 Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #58963 [ run ] completed with state ABORTED. Commit: 18908a5

Link to invocation

…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>
@luyiyun1021 luyiyun1021 force-pushed the fix/visualgen-spawn-response-handle branch from c0e0c33 to 9edf869 Compare July 13, 2026 12:17
@luyiyun1021 luyiyun1021 changed the title [TRTLLM-14040][fix] VisualGen: declare same-process client at the launch site [TRTLLM-14040][fix] Fix multi-GPU VisualGen response IPC: REBUILD_LOCAL handle crossed a process boundary Jul 13, 2026
@luyiyun1021

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@luyiyun1021 luyiyun1021 requested a review from zhenhuaw-me July 13, 2026 12:21
@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #58970 [ run ] triggered by Bot. Commit: 9edf869 Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #58967 [ run ] completed with state ABORTED. Commit: c0e0c33

Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #58970 [ run ] completed with state SUCCESS. Commit: 9edf869
/LLM/main/L0_MergeRequest_PR pipeline #47502 completed with status: 'SUCCESS'

CI Report

Link to invocation

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>
@luyiyun1021 luyiyun1021 force-pushed the fix/visualgen-spawn-response-handle branch from 9edf869 to d676a01 Compare July 14, 2026 03:21
@luyiyun1021

Copy link
Copy Markdown
Collaborator Author

/bot kill

@luyiyun1021

Copy link
Copy Markdown
Collaborator Author

/bot reuse-pipeline

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #59087 [ reuse-pipeline ] triggered by Bot. Commit: d676a01 Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #59089 [ kill ] triggered by Bot. Commit: d676a01 Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #59087 [ reuse-pipeline ] completed with state ABORTED. Commit: d676a01
Can't reuse PR_Github #0 with status: UNKNOWN

Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #59089 [ kill ] completed with state SUCCESS. Commit: d676a01
Successfully killed previous jobs for commit d676a01

Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #59090 [ run ] triggered by Bot. Commit: d676a01 Link to invocation

@luyiyun1021

Copy link
Copy Markdown
Collaborator Author

/bot reuse-pipeline --number 58970

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #59091 [ reuse-pipeline ] triggered by Bot. Commit: d676a01 Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #59090 [ run ] completed with state ABORTED. Commit: d676a01

Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #59091 [ reuse-pipeline ] completed with state SUCCESS. Commit: d676a01
Reusing PR_Github #59090 for commit d676a01

Link to invocation

@luyiyun1021 luyiyun1021 merged commit 67e8dc9 into NVIDIA:main Jul 14, 2026
7 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants