[https://nvbugs/6084743][fix] Use free port for FLUX/WAN multi-GPU test distributed init

[karljang]

Summary

• Fixes the flaky post-merge failure test_flux_pipeline.py::TestFluxCombinedOptimizations::test_all_optimizations_combined tracked by nvbug 6084743 / TRTLLM-12007.
• _setup_distributed in tests/unittest/_torch/visual_gen/test_flux_pipeline.py hardcoded MASTER_PORT=12355 for torch.distributed init. When two multi-GPU tests land on the same node back-to-back, or another process already holds 12355, init_process_group fails with "The server socket has failed to listen on any local network address".
• Thread master_port through _setup_distributed and both worker functions, and allocate a fresh free port via tensorrt_llm._utils.get_free_port() at each mp.spawn call site.
• Same pattern is already used by the sibling tests under tests/unittest/_torch/visual_gen/multi_gpu/ (e.g. test_flux_ulysses.py, test_ulysses_attention.py).

Test plan

• python -m py_compile tests/unittest/_torch/visual_gen/test_flux_pipeline.py
• Local test pass on B200 x2 (pre-existing flake doesn't reproduce with fixed port collision path).
• CI: /bot run

Summary by CodeRabbit

Tests

• Improved distributed multi-GPU test infrastructure with dynamic port allocation
• Tests can now run reliably in parallel without port conflicts
• Enhanced stability across diverse system environments and resource configurations
• Better support for containerized and shared testing infrastructure

[karljang]

/bot run --disble-fail-fast

[coderabbitai]

📝 Walkthrough

Walkthrough

The change adds dynamic TCP port selection for distributed multi-GPU test workers by introducing a master_port parameter, threading it through internal test helper functions, and replacing hardcoded port values with get_free_port() allocation.

Changes

Cohort / File(s)

Summary

Dynamic Port Selection for Distributed Tests tests/unittest/_torch/visual_gen/test_flux_pipeline.py

Imports get_free_port, adds master_port parameter to _setup_distributed, _run_ulysses_worker, and _run_all_optimizations_worker functions. Updates test entrypoints test_ulysses_2gpu_correctness and test_all_optimizations_combined to call get_free_port() and pass the dynamic port value through the worker spawn calls. Modifies distributed initialization to use the provided master_port instead of a hardcoded constant.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~10 minutes

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.
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.
Description check	✅ Passed	The pull request description is comprehensive and follows the template structure with clear summary, test coverage, and checklist items completed.
Title check	✅ Passed	The title accurately describes the main change: using dynamic free port allocation instead of hardcoded port for multi-GPU distributed tests.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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

[coderabbitai]

🧹 Nitpick comments (2)

tests/unittest/_torch/visual_gen/test_flux_pipeline.py (2)

974-975: QA list update is unnecessary for this PR.

This is a unittest-only fix under tests/unittest/...; no tests/integration/test_lists/qa/ changes are needed.

As per coding guidelines: "If the PR only touches unittest/ or narrow unit scope, say explicitly whether QA list updates are unnecessary or optional."

Also applies to: 1157-1158

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@tests/unittest/_torch/visual_gen/test_flux_pipeline.py` around lines 974 -
975, The change only touches unit tests (function test_ulysses_2gpu_correctness
in tests/unittest/_torch/visual_gen/test_flux_pipeline.py), so remove or revert
any edits to the QA lists under tests/integration/test_lists/qa/ and update the
PR description to explicitly state that QA list updates are unnecessary for this
unittest-only PR; ensure the commit message or PR body references the affected
unit test (test_ulysses_2gpu_correctness) and clarifies no integration QA
changes are required.

---

1028-1034: Optional hardening: add a bounded retry around spawn for residual port race windows.

At Line 1028 and Line 1212, the port is selected in the parent process before children bind; another process can still claim it in between. A small retry loop would make this even less flaky.

💡 Suggested hardening diff

+def _spawn_with_retry(worker_fn, world_size, checkpoint_path, inputs_cpu, return_dict, max_retries=3):
+    for attempt in range(max_retries):
+        master_port = get_free_port()
+        try:
+            mp.spawn(
+                worker_fn,
+                args=(world_size, master_port, checkpoint_path, inputs_cpu, return_dict),
+                nprocs=world_size,
+                join=True,
+            )
+            return
+        except RuntimeError as err:
+            msg = str(err)
+            if attempt + 1 < max_retries and (
+                "failed to listen on any local network address" in msg
+                or "Address already in use" in msg
+            ):
+                continue
+            raise
@@
-        master_port = get_free_port()
-        mp.spawn(
-            _run_ulysses_worker,
-            args=(2, master_port, FLUX1_CHECKPOINT_PATH, inputs_cpu, return_dict),
-            nprocs=2,
-            join=True,
-        )
+        _spawn_with_retry(
+            _run_ulysses_worker, 2, FLUX1_CHECKPOINT_PATH, inputs_cpu, return_dict
+        )
@@
-        master_port = get_free_port()
-        mp.spawn(
-            _run_all_optimizations_worker,
-            args=(2, master_port, FLUX1_CHECKPOINT_PATH, inputs_cpu, return_dict),
-            nprocs=2,
-            join=True,
-        )
+        _spawn_with_retry(
+            _run_all_optimizations_worker, 2, FLUX1_CHECKPOINT_PATH, inputs_cpu, return_dict
+        )

Also applies to: 1212-1218

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@tests/unittest/_torch/visual_gen/test_flux_pipeline.py` around lines 1028 -
1034, The test selects master_port via get_free_port() then immediately calls
mp.spawn(_run_ulysses_worker...), which can race if another process grabs the
port; wrap the get_free_port() + mp.spawn call in a small bounded retry loop
(e.g., try up to N times with a short sleep) that re-reads master_port and
re-attempts mp.spawn when spawn fails or when a quick port-binding check fails;
apply the same pattern to the other occurrence around lines using
master_port/get_free_port() and mp.spawn so both spots (the callsites of
get_free_port(), mp.spawn, and the target function _run_ulysses_worker) are
hardened against transient port races.

🤖 Prompt for all review comments with AI agents

Verify each finding against the current code and only fix it if needed.

Nitpick comments:
In `@tests/unittest/_torch/visual_gen/test_flux_pipeline.py`:
- Around line 974-975: The change only touches unit tests (function
test_ulysses_2gpu_correctness in
tests/unittest/_torch/visual_gen/test_flux_pipeline.py), so remove or revert any
edits to the QA lists under tests/integration/test_lists/qa/ and update the PR
description to explicitly state that QA list updates are unnecessary for this
unittest-only PR; ensure the commit message or PR body references the affected
unit test (test_ulysses_2gpu_correctness) and clarifies no integration QA
changes are required.
- Around line 1028-1034: The test selects master_port via get_free_port() then
immediately calls mp.spawn(_run_ulysses_worker...), which can race if another
process grabs the port; wrap the get_free_port() + mp.spawn call in a small
bounded retry loop (e.g., try up to N times with a short sleep) that re-reads
master_port and re-attempts mp.spawn when spawn fails or when a quick
port-binding check fails; apply the same pattern to the other occurrence around
lines using master_port/get_free_port() and mp.spawn so both spots (the
callsites of get_free_port(), mp.spawn, and the target function
_run_ulysses_worker) are hardened against transient port races.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro Plus

Run ID: f46c5caf-3f42-44ed-a22c-2499fd75f526

📥 Commits

Reviewing files that changed from the base of the PR and between 7b0f60e and 111dda9.

📒 Files selected for processing (1)

• tests/unittest/_torch/visual_gen/test_flux_pipeline.py

[karljang]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #45123 Bot args parsing error: usage: /bot [-h]
{run,kill,skip,submit,reviewers,reuse-pipeline,reuse-review} ...
/bot: error: unrecognized arguments: --disble-fail-fast

Link to invocation

[tensorrt-cicd]

PR_Github #45124 [ run ] triggered by Bot. Commit: 111dda9 Link to invocation

[tensorrt-cicd]

PR_Github #45124 [ run ] completed with state SUCCESS. Commit: 111dda9
/LLM/main/L0_MergeRequest_PR pipeline #35417 completed with status: 'SUCCESS'

CI Report

Link to invocation

[karljang]

@chang-l , a few unit tests have been fixed by this PR, could you please review? :)