From 2a8c5d20b1e85040e463730cbeaef0355c18466b Mon Sep 17 00:00:00 2001 From: Ying Guo <244492186+yingguo-trt@users.noreply.github.com> Date: Wed, 24 Jun 2026 17:49:23 +0800 Subject: [PATCH 01/22] [TRTLLM-13304][test] Add VisualGen multi-node QA coverage Add QA coverage for multi-node VisualGen (dev epic TRTLLM-11407): - SLURM 2-node x 8-GPU end-to-end LPIPS test for cfg2_attn2d_2x2_ulysses2, exercising the real _detect_external_launch SLURM path; parent pre-checks the checkpoint and validates srun output to avoid an all-skip false pass. - world_size=16 rank/group unit tests (attn2d and ring) using a fake DeviceMesh, extending coverage beyond the 8-GPU single-node lane. - Register the multi-node case in qa/llm_function_multinode.txt. Signed-off-by: Ying Guo <244492186+yingguo-trt@users.noreply.github.com> --- .../visual_gen/test_visual_gen_multi_gpu.py | 268 ++++++++++++++++++ .../test_lists/qa/llm_function_multinode.txt | 1 + .../multi_gpu/test_visual_gen_mapping.py | 179 ++++++++++++ 3 files changed, 448 insertions(+) diff --git a/tests/integration/defs/examples/visual_gen/test_visual_gen_multi_gpu.py b/tests/integration/defs/examples/visual_gen/test_visual_gen_multi_gpu.py index 13f15e464b09..dd8e99477e67 100644 --- a/tests/integration/defs/examples/visual_gen/test_visual_gen_multi_gpu.py +++ b/tests/integration/defs/examples/visual_gen/test_visual_gen_multi_gpu.py @@ -16,7 +16,10 @@ import glob import os +import shutil +import subprocess import sys +from pathlib import Path from typing import Callable import pytest @@ -39,6 +42,7 @@ _run_lpips_eval, _run_wan_lpips_pipeline, _save_lpips_video_mp4, + _skip_if_missing, ) @@ -72,6 +76,17 @@ def _parallel_config(**kwargs): ("tp2_attn2d_2x1", {"tp_size": 2, "attn2d_size": (2, 1)}), ] +WAN22_LPIPS_MULTINODE_WORLD_SIZE = 16 +WAN22_LPIPS_MULTINODE_NODES = 2 +WAN22_LPIPS_MULTINODE_GPUS_PER_NODE = 8 +WAN22_LPIPS_MULTINODE_VARIANTS = [ + ( + "cfg2_attn2d_2x2_ulysses2", + {"cfg_size": 2, "attn2d_size": (2, 2), "ulysses_size": 2}, + ), +] +_MULTINODE_SLURM_CHILD_ENV = "TRTLLM_VISUAL_GEN_MULTINODE_SLURM_CHILD" + @pytest.fixture(autouse=True, scope="module") def _cleanup_mpi_env(): @@ -175,6 +190,243 @@ def _skip_if_insufficient_gpus_for_parallel(parallel): ) +def _slurm_rank_env(): + if "SLURM_PROCID" not in os.environ or "SLURM_NTASKS" not in os.environ: + return None + return int(os.environ["SLURM_PROCID"]), int(os.environ["SLURM_NTASKS"]) + + +def _default_master_port(): + job_id = int(os.environ.get("SLURM_JOB_ID", "0") or 0) + return str(20000 + job_id % 20000) + + +def _slurm_node_count(): + for var in ("SLURM_JOB_NUM_NODES", "SLURM_NNODES"): + if var in os.environ: + return int(os.environ[var]) + return None + + +def _trtllm_launch_wrapper_world_size(): + try: + return int(os.environ.get("tllm_mpi_size", "1") or 1) + except ValueError: + return 1 + + +def _multinode_subprocess_timeout(): + return int(os.environ.get("TRTLLM_VISUAL_GEN_MULTINODE_TIMEOUT", "3600")) + + +def _resolve_slurm_master_addr(): + if os.environ.get("MASTER_ADDR"): + return os.environ["MASTER_ADDR"] + + nodelist = os.environ.get("SLURM_JOB_NODELIST") + if not nodelist: + pytest.skip("SLURM_JOB_NODELIST is required to resolve MASTER_ADDR") + if shutil.which("scontrol") is None: + pytest.skip("scontrol is required to resolve MASTER_ADDR from SLURM_JOB_NODELIST") + + result = subprocess.run( + ["scontrol", "show", "hostnames", nodelist], + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + check=False, + ) + if result.returncode != 0: + pytest.fail(f"Failed to resolve SLURM master host:\n{result.stdout}") + + master_addr = next((line.strip() for line in result.stdout.splitlines() if line.strip()), "") + if not master_addr: + pytest.fail(f"scontrol returned no hostnames for SLURM_JOB_NODELIST={nodelist!r}") + os.environ["MASTER_ADDR"] = master_addr + return master_addr + + +def _ensure_slurm_external_launch_env(): + os.environ["MASTER_ADDR"] = _resolve_slurm_master_addr() + os.environ.setdefault("MASTER_PORT", _default_master_port()) + + # Force VisualGen to exercise its SLURM detection branch even when the + # surrounding CI wrapper leaves torchrun-like variables behind. + for var in ("RANK", "WORLD_SIZE", "LOCAL_RANK"): + os.environ.pop(var, None) + + +def _run_wan22_multinode_slurm_rank(tmp_path, variant_name, parallel): + if not MODULES_AVAILABLE: + pytest.skip("Required modules not available") + + rank_env = _slurm_rank_env() + if rank_env is None: + pytest.skip("This VisualGen multi-node case must run under SLURM rank env") + rank, world_size = rank_env + if world_size != WAN22_LPIPS_MULTINODE_WORLD_SIZE: + pytest.skip(f"Requires {WAN22_LPIPS_MULTINODE_WORLD_SIZE} SLURM tasks, got {world_size}") + node_count = _slurm_node_count() + if node_count is not None and node_count < WAN22_LPIPS_MULTINODE_NODES: + pytest.skip( + f"Requires at least {WAN22_LPIPS_MULTINODE_NODES} SLURM nodes, got {node_count}" + ) + + _ensure_slurm_external_launch_env() + ParallelConfig(**parallel).validate_world_size(world_size) + model_path = _lpips_model_path("Wan2.2-T2V-A14B-Diffusers") + _skip_if_missing(model_path, "Wan 2.2 checkpoint", is_dir=True) + + from tensorrt_llm import VisualGen, VisualGenArgs, VisualGenParams + from tensorrt_llm.visual_gen.args import AttentionConfig, CompilationConfig, TorchCompileConfig + + visual_gen_args = VisualGenArgs( + model=model_path, + compilation_config=CompilationConfig(skip_warmup=True), + torch_compile_config=TorchCompileConfig(enable=False), + attention_config=AttentionConfig(backend="FA4"), + parallel_config=parallel, + ) + + visual_gen = None + try: + try: + visual_gen = VisualGen(model=model_path, args=visual_gen_args) + except SystemExit as exc: + assert rank != 0, "Only non-zero SLURM ranks should exit through worker mode" + assert exc.code in (0, None) + return + + assert rank == 0 + params = VisualGenParams( + height=WAN22_LPIPS_HEIGHT, + width=WAN22_LPIPS_WIDTH, + num_frames=WAN22_LPIPS_NUM_FRAMES, + num_inference_steps=WAN22_LPIPS_NUM_INFERENCE_STEPS, + guidance_scale=WAN22_LPIPS_GUIDANCE_SCALE, + seed=WAN22_LPIPS_SEED, + frame_rate=WAN22_LPIPS_FRAME_RATE, + negative_prompt=WAN22_LPIPS_NEGATIVE_PROMPT, + ) + output = visual_gen.generate(inputs=WAN22_LPIPS_PROMPT, params=params) + assert output.error is None, f"unexpected error on Wan 2.2 multi-node run: {output.error}" + assert output.video is not None + + generated_path = tmp_path / f"wan22_t2v_generated_{variant_name}_slurm.mp4" + output.save(generated_path, frame_rate=WAN22_LPIPS_FRAME_RATE) + assert generated_path.is_file(), ( + f"VisualGen multi-node run did not produce {generated_path}" + ) + + golden_path = _golden_media_path( + tmp_path, "wan22_t2v_lpips_golden_video.mp4", "Wan 2.2 LPIPS golden video" + ) + score = _run_lpips_eval( + tmp_path, + f"wan22_t2v_{variant_name}_slurm", + "video", + WAN22_LPIPS_PROMPT, + golden_path, + generated_path, + ) + _assert_lpips_below_threshold(score, WAN_MULTI_GPU_LPIPS_THRESHOLD) + finally: + if visual_gen is not None: + visual_gen.shutdown() + + +def _run_wan22_multinode_slurm_parent(variant_name): + if ( + os.environ.get("TLLM_SPAWN_PROXY_PROCESS") == "1" + and _trtllm_launch_wrapper_world_size() >= WAN22_LPIPS_MULTINODE_WORLD_SIZE + ): + pytest.fail( + "VisualGen SLURM external-launch coverage cannot run under " + "trtllm-llmapi-launch because that wrapper removes SLURM_* env " + "before user code. Run this nodeid with direct srun so " + "_detect_external_launch() sees the real SLURM rank environment." + ) + + if os.environ.get("SLURM_JOB_ID") is None: + pytest.skip("A SLURM allocation is required for the VisualGen multi-node LPIPS case") + node_count = _slurm_node_count() + if node_count is not None and node_count < WAN22_LPIPS_MULTINODE_NODES: + pytest.skip( + f"Requires at least {WAN22_LPIPS_MULTINODE_NODES} SLURM nodes, got {node_count}" + ) + if shutil.which("srun") is None: + pytest.skip("srun is required for the VisualGen multi-node LPIPS case") + + # Pre-check the checkpoint here so a missing model skips the parent honestly, + # instead of letting every srun rank skip and the parent report a false pass + # (an all-skipped pytest run still exits 0). + _skip_if_missing( + _lpips_model_path("Wan2.2-T2V-A14B-Diffusers"), + "Wan 2.2 checkpoint", + is_dir=True, + ) + + env = os.environ.copy() + env[_MULTINODE_SLURM_CHILD_ENV] = "1" + env["MASTER_ADDR"] = _resolve_slurm_master_addr() + env.setdefault("MASTER_PORT", _default_master_port()) + env["PYTHONUNBUFFERED"] = "1" + for var in ("RANK", "WORLD_SIZE", "LOCAL_RANK"): + env.pop(var, None) + + test_file = str(Path(__file__).resolve()) + nodeid = f"{test_file}::test_wan22_t2v_lpips_against_golden_multinode_slurm[{variant_name}]" + cmd = [ + "srun", + "-l", + "--overlap", + f"--nodes={WAN22_LPIPS_MULTINODE_NODES}", + f"--ntasks={WAN22_LPIPS_MULTINODE_WORLD_SIZE}", + f"--ntasks-per-node={WAN22_LPIPS_MULTINODE_GPUS_PER_NODE}", + "--export=ALL", + sys.executable, + "-m", + "pytest", + "-q", + "-s", + nodeid, + ] + try: + result = subprocess.run( + cmd, + cwd=Path(__file__).resolve().parents[5], + env=env, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + check=False, + timeout=_multinode_subprocess_timeout(), + ) + except subprocess.TimeoutExpired as exc: + output = exc.output or "" + pytest.fail( + "VisualGen multi-node SLURM subprocess timed out after " + f"{exc.timeout} seconds:\n{output}" + ) + + if result.returncode != 0: + pytest.fail( + "VisualGen multi-node SLURM subprocess failed with " + f"exit code {result.returncode}:\n{result.stdout}" + ) + + # Guard against a false pass: an all-skipped/empty pytest run also exits 0. + # The parent has already validated every precondition, so once srun launches + # all ranks must actually pass; any skip or empty collection is a real defect. + output = result.stdout or "" + if "passed" not in output or "skipped" in output or "no tests ran" in output: + pytest.fail( + "VisualGen multi-node SLURM run exited 0 but did not actually execute " + "(expected all ranks to pass with no skips):\n" + f"{output}" + ) + + def _wan22_lpips_distributed_worker(rank: int, world_size: int, **kwargs) -> None: parallel = kwargs["parallel"] _parallel_config(**parallel).validate_world_size(world_size) @@ -265,3 +517,19 @@ def test_wan22_t2v_lpips_against_golden_multi_gpu( ) def test_wan22_t2v_lpips_against_golden_tp(_visual_gen_deps, tmp_path, variant_name, parallel): _run_wan22_t2v_lpips_case(tmp_path, variant_name, parallel) + + +@pytest.mark.parametrize( + "variant_name,parallel", + WAN22_LPIPS_MULTINODE_VARIANTS, + ids=[name for name, _ in WAN22_LPIPS_MULTINODE_VARIANTS], +) +def test_wan22_t2v_lpips_against_golden_multinode_slurm(tmp_path, variant_name, parallel): + if _slurm_rank_env() is not None: + _run_wan22_multinode_slurm_rank(tmp_path, variant_name, parallel) + return + + if os.environ.get(_MULTINODE_SLURM_CHILD_ENV): + pytest.skip("VisualGen SLURM child was not launched with SLURM rank env") + + _run_wan22_multinode_slurm_parent(variant_name) diff --git a/tests/integration/test_lists/qa/llm_function_multinode.txt b/tests/integration/test_lists/qa/llm_function_multinode.txt index cb10040d1934..10ca417883d9 100644 --- a/tests/integration/test_lists/qa/llm_function_multinode.txt +++ b/tests/integration/test_lists/qa/llm_function_multinode.txt @@ -7,3 +7,4 @@ test_e2e.py::test_multi_nodes_eval[MiniMax-M3-tp16-mmlu] test_e2e.py::test_openai_disagg_multi_nodes_completion[ctx_tp2pp1-gen_tp2pp1] test_e2e.py::test_openai_disagg_multi_nodes_completion[ctx_tp1pp2-gen_tp1pp2] test_e2e.py::test_openai_disagg_multi_nodes_completion_service_discovery[etcd] +examples/visual_gen/test_visual_gen_multi_gpu.py::test_wan22_t2v_lpips_against_golden_multinode_slurm[cfg2_attn2d_2x2_ulysses2] diff --git a/tests/unittest/_torch/visual_gen/multi_gpu/test_visual_gen_mapping.py b/tests/unittest/_torch/visual_gen/multi_gpu/test_visual_gen_mapping.py index c73762ae3c7a..725b0bf7aff1 100644 --- a/tests/unittest/_torch/visual_gen/multi_gpu/test_visual_gen_mapping.py +++ b/tests/unittest/_torch/visual_gen/multi_gpu/test_visual_gen_mapping.py @@ -72,6 +72,68 @@ def _run_multi_gpu(world_size, test_fn): ) +def _row_major_coords(rank, dim_names, dim_sizes): + # Mirrors torch.distributed.device_mesh.init_device_mesh row-major layout; + # the existing multi-GPU mapping tests anchor this assumption against real DeviceMesh. + coords = {} + remaining = rank + for dim in reversed(dim_names): + size = dim_sizes[dim] + coords[dim] = remaining % size + remaining //= size + return coords + + +class _FakeDeviceMeshView: + def __init__(self, dim_names, dim_sizes, rank, group_dims): + self._dim_names = dim_names + self._dim_sizes = dim_sizes + self._rank = rank + self._group_dims = tuple(group_dims) + self._coords = _row_major_coords(rank, dim_names, dim_sizes) + self._world_size = 1 + for size in dim_sizes.values(): + self._world_size *= size + + def __getitem__(self, dim): + if isinstance(dim, tuple): + group_dims = dim + else: + group_dims = (dim,) + return _FakeDeviceMeshView( + self._dim_names, + self._dim_sizes, + self._rank, + group_dims, + ) + + def _flatten(self, mesh_dim_name): + return _FakeDeviceMeshView( + self._dim_names, + self._dim_sizes, + self._rank, + self._group_dims, + ) + + def get_local_rank(self): + assert len(self._group_dims) == 1 + return self._coords[self._group_dims[0]] + + def get_group(self): + fixed_dims = set(self._dim_names) - set(self._group_dims) + group = [] + for rank in range(self._world_size): + coords = _row_major_coords(rank, self._dim_names, self._dim_sizes) + if all(coords[dim] == self._coords[dim] for dim in fixed_dims): + group.append(rank) + return tuple(group) + + +class _FakeDeviceMesh(_FakeDeviceMeshView): + def __init__(self, dim_names, dim_sizes, rank): + super().__init__(dim_names, dim_sizes, rank, dim_names) + + # ============================================================================= # Single-GPU tests (no dist required) # ============================================================================= @@ -273,6 +335,123 @@ def test_mixed_parallelism(self): assert m.world_size == 2 +class TestRankLinearizationWithoutDist: + def test_world_size_16_cfg_attn2d_ulysses_rank_groups(self): + from tensorrt_llm._torch.device_mesh import DeviceMeshTopologyImpl + + old_mesh = DeviceMeshTopologyImpl.device_mesh + old_seq_mesh = VisualGenMapping.seq_mesh + try: + for rank in range(16): + DeviceMeshTopologyImpl.device_mesh = None + VisualGenMapping.seq_mesh = None + vgm = VisualGenMapping( + world_size=16, + rank=rank, + cfg_size=2, + attn2d_row_size=2, + attn2d_col_size=2, + ulysses_size=2, + ) + + fake_mesh = _FakeDeviceMesh(vgm._dim_names, vgm._dim_sizes, rank) + DeviceMeshTopologyImpl.device_mesh = fake_mesh + VisualGenMapping.seq_mesh = fake_mesh["cp_row", "cp_col", "ulysses"]._flatten( + mesh_dim_name="seq" + ) + vgm._attach_attn2d_groups_from_device_mesh() + + expected_cfg_rank = rank // 8 + expected_cp_row_rank = (rank // 4) % 2 + expected_cp_col_rank = (rank // 2) % 2 + expected_ulysses_rank = rank % 2 + expected_cp_rank = expected_cp_row_rank * 2 + expected_cp_col_rank + expected_seq_rank = expected_cp_rank * 2 + expected_ulysses_rank + + assert vgm.cfg_rank == expected_cfg_rank + assert vgm.tp_rank == 0 + assert vgm.cp_row_rank == expected_cp_row_rank + assert vgm.cp_col_rank == expected_cp_col_rank + assert vgm.cp_rank == expected_cp_rank + assert vgm.attn2d_mesh_rank == expected_cp_rank + assert vgm.ulysses_rank == expected_ulysses_rank + assert vgm.seq_rank == expected_seq_rank + assert vgm.seq_size == 8 + assert vgm.is_cfg_conditional == (rank < 8) + + cfg_pair_start = rank % 8 + ulysses_pair_start = rank - expected_ulysses_rank + seq_group_start = expected_cfg_rank * 8 + cp_group_start = expected_cfg_rank * 8 + expected_ulysses_rank + row_group_start = ( + expected_cfg_rank * 8 + expected_cp_row_rank * 4 + expected_ulysses_rank + ) + col_group_start = ( + expected_cfg_rank * 8 + expected_cp_col_rank * 2 + expected_ulysses_rank + ) + + assert vgm.cfg_group == (cfg_pair_start, cfg_pair_start + 8) + assert vgm.ulysses_group == (ulysses_pair_start, ulysses_pair_start + 1) + assert vgm.cp_group == tuple(cp_group_start + stride for stride in (0, 2, 4, 6)) + assert vgm.attn2d_mesh_group == vgm.cp_group + assert vgm.attn2d_row_group == (row_group_start, row_group_start + 2) + assert vgm.attn2d_col_group == (col_group_start, col_group_start + 4) + assert vgm.seq_group() == tuple(range(seq_group_start, seq_group_start + 8)) + finally: + DeviceMeshTopologyImpl.device_mesh = old_mesh + VisualGenMapping.seq_mesh = old_seq_mesh + + def test_world_size_16_cfg_ring_ulysses_rank_groups(self): + from tensorrt_llm._torch.device_mesh import DeviceMeshTopologyImpl + + old_mesh = DeviceMeshTopologyImpl.device_mesh + old_seq_mesh = VisualGenMapping.seq_mesh + try: + for rank in range(16): + DeviceMeshTopologyImpl.device_mesh = None + VisualGenMapping.seq_mesh = None + vgm = VisualGenMapping( + world_size=16, + rank=rank, + cfg_size=2, + ring_size=4, + ulysses_size=2, + ) + + fake_mesh = _FakeDeviceMesh(vgm._dim_names, vgm._dim_sizes, rank) + DeviceMeshTopologyImpl.device_mesh = fake_mesh + VisualGenMapping.seq_mesh = fake_mesh["cp", "ulysses"]._flatten(mesh_dim_name="seq") + + expected_cfg_rank = rank // 8 + expected_cp_rank = (rank // 2) % 4 + expected_ulysses_rank = rank % 2 + expected_seq_rank = expected_cp_rank * 2 + expected_ulysses_rank + + assert vgm.cfg_rank == expected_cfg_rank + assert vgm.tp_rank == 0 + assert vgm.cp_rank == expected_cp_rank + assert vgm.ring_rank == expected_cp_rank + assert vgm.ulysses_rank == expected_ulysses_rank + assert vgm.seq_rank == expected_seq_rank + assert vgm.seq_size == 8 + assert vgm.is_cfg_conditional == (rank < 8) + + cfg_pair_start = rank % 8 + ulysses_pair_start = rank - expected_ulysses_rank + seq_group_start = expected_cfg_rank * 8 + cp_group_start = expected_cfg_rank * 8 + expected_ulysses_rank + + assert vgm.cfg_group == (cfg_pair_start, cfg_pair_start + 8) + assert vgm.tp_group_pg == (rank,) + assert vgm.ulysses_group == (ulysses_pair_start, ulysses_pair_start + 1) + assert vgm.cp_group == tuple(cp_group_start + stride for stride in (0, 2, 4, 6)) + assert vgm.ring_group == vgm.cp_group + assert vgm.seq_group() == tuple(range(seq_group_start, seq_group_start + 8)) + finally: + DeviceMeshTopologyImpl.device_mesh = old_mesh + VisualGenMapping.seq_mesh = old_seq_mesh + + # ============================================================================= # Multi-GPU tests — validate actual DeviceMesh groups and ranks # ============================================================================= From 51c20ae610d400a4c66505be9d96f8c03095c487 Mon Sep 17 00:00:00 2001 From: Ying Guo <244492186+yingguo-trt@users.noreply.github.com> Date: Thu, 25 Jun 2026 14:57:34 +0800 Subject: [PATCH 02/22] [TRTLLM-13304][test] Skip VisualGen multinode SLURM case under trtllm-llmapi-launch The wan22 multinode SLURM case needs real SLURM_* env for rank detection, but trtllm-llmapi-launch strips SLURM_* before user code. Downgrade the hard pytest.fail to pytest.skip so the case skips gracefully under that launcher, consistent with its other environment prerequisite checks, instead of failing the pipeline. Signed-off-by: Ying Guo <244492186+yingguo-trt@users.noreply.github.com> --- .../defs/examples/visual_gen/test_visual_gen_multi_gpu.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/integration/defs/examples/visual_gen/test_visual_gen_multi_gpu.py b/tests/integration/defs/examples/visual_gen/test_visual_gen_multi_gpu.py index dd8e99477e67..0c30df8d2398 100644 --- a/tests/integration/defs/examples/visual_gen/test_visual_gen_multi_gpu.py +++ b/tests/integration/defs/examples/visual_gen/test_visual_gen_multi_gpu.py @@ -340,8 +340,8 @@ def _run_wan22_multinode_slurm_parent(variant_name): os.environ.get("TLLM_SPAWN_PROXY_PROCESS") == "1" and _trtllm_launch_wrapper_world_size() >= WAN22_LPIPS_MULTINODE_WORLD_SIZE ): - pytest.fail( - "VisualGen SLURM external-launch coverage cannot run under " + pytest.skip( + "VisualGen SLURM external-launch coverage is skipped under " "trtllm-llmapi-launch because that wrapper removes SLURM_* env " "before user code. Run this nodeid with direct srun so " "_detect_external_launch() sees the real SLURM rank environment." From a2b802a9e595f707825114976087841297637b51 Mon Sep 17 00:00:00 2001 From: Ying Guo <244492186+yingguo-trt@users.noreply.github.com> Date: Thu, 25 Jun 2026 17:36:55 +0800 Subject: [PATCH 03/22] [TRTLLM-13304][ci] Add dedicated L0 stage for VisualGen multinode SLURM test VisualGen multinode uses torch.distributed + NCCL (TLLM_DISABLE_MPI=1) and must NOT be wrapped with trtllm-llmapi-launch, which strips SLURM_* env. Add a dedicated l0_b200_visual_gen_multinode list + 2-node/16-GPU post-merge stage that launches without the MPI wrapper: one task per node sets up, then rank 0 runs the parent pytest with SLURM rank env stripped so the test self-orchestrates its own 16-rank srun. WIP pending dev confirmation of the intended VisualGen multinode launch recipe (TRTLLM-11407 still in progress). Signed-off-by: Ying Guo <244492186+yingguo-trt@users.noreply.github.com> --- jenkins/L0_MergeRequest.groovy | 1 + jenkins/L0_Test.groovy | 19 +++++++++-- jenkins/scripts/slurm_run.sh | 20 +++++++++++ .../visual_gen/test_visual_gen_multi_gpu.py | 4 +-- .../test-db/l0_b200_visual_gen_multinode.yml | 33 +++++++++++++++++++ 5 files changed, 73 insertions(+), 4 deletions(-) create mode 100644 tests/integration/test_lists/test-db/l0_b200_visual_gen_multinode.yml diff --git a/jenkins/L0_MergeRequest.groovy b/jenkins/L0_MergeRequest.groovy index 981b9dc403c7..155269180596 100644 --- a/jenkins/L0_MergeRequest.groovy +++ b/jenkins/L0_MergeRequest.groovy @@ -1109,6 +1109,7 @@ def getMultiGpuFileChanged(pipeline, testFilter, globalVars) "tests/integration/defs/cpp/test_multi_gpu.py", "tests/integration/test_lists/test-db/l0_b200_multi_gpus_perf_sanity.yml", "tests/integration/test_lists/test-db/l0_b200_multi_nodes_perf_sanity_ctx1_node1_gpu4_gen1_node1_gpu8.yml", + "tests/integration/test_lists/test-db/l0_b200_visual_gen_multinode.yml", "tests/integration/test_lists/test-db/l0_b200_visual_gen_perf_sanity.yml", "tests/integration/test_lists/test-db/l0_dgx_b200.yml", "tests/integration/test_lists/test-db/l0_dgx_b300.yml", diff --git a/jenkins/L0_Test.groovy b/jenkins/L0_Test.groovy index 51c6e3d73575..7cdf229f6ea4 100644 --- a/jenkins/L0_Test.groovy +++ b/jenkins/L0_Test.groovy @@ -1362,6 +1362,16 @@ def getNodeArgs(int nodeCount, int gpuCount, boolean setSegment = false) { return args } +def getVisualGenMultinodeParentNodeArgs(int nodeCount, int gpuCount) { + int gpusPerNode = ((gpuCount / nodeCount) as BigDecimal).setScale(0, BigDecimal.ROUND_CEILING).intValue() + return [ + "--nodes=${nodeCount}", + "--ntasks=${nodeCount}", + "--ntasks-per-node=1", + "--gpus-per-node=${gpusPerNode}", + ] +} + def getPytestBaseCommandLine( String llmSrc, String stageName, @@ -1671,7 +1681,8 @@ def runLLMTestlistWithSbatch(pipeline, platform, testList, config=VANILLA_CONFIG // Generate Pytest command String pytestUtil = "" - if (nodeCount > 1) { + def visualGenMultinodeSlurmMode = testList == "l0_b200_visual_gen_multinode" + if (nodeCount > 1 && !visualGenMultinodeSlurmMode) { pytestUtil = "$llmSrcNode/tensorrt_llm/llmapi/trtllm-llmapi-launch" } def uploadPath = "${env.JOB_NAME}/${env.BUILD_NUMBER}" @@ -1715,7 +1726,9 @@ def runLLMTestlistWithSbatch(pipeline, platform, testList, config=VANILLA_CONFIG // Generate Job Launch Script def container = LLM_DOCKER_IMAGE.replace("urm.nvidia.com/", "urm.nvidia.com#") def mounts = getMountListForSlurmTest(cluster, true).join(",") - String[] taskArgs = getNodeArgs(nodeCount, gpuCount, disaggMultiNodeMode) + String[] taskArgs = visualGenMultinodeSlurmMode ? + getVisualGenMultinodeParentNodeArgs(nodeCount, gpuCount) : + getNodeArgs(nodeCount, gpuCount, disaggMultiNodeMode) if (taskArgs == null) { error "Invalid Slurm test stage name is set" } @@ -1875,6 +1888,7 @@ def runLLMTestlistWithSbatch(pipeline, platform, testList, config=VANILLA_CONFIG export resourcePathNode=$resourcePathNode export pytestCommand="$pytestCommand" export coverageConfigFile="$coverageConfigFile" + export TRTLLM_VISUAL_GEN_MULTINODE_SLURM_PARENT=${visualGenMultinodeSlurmMode ? "1" : "0"} export HF_TOKEN=$HF_TOKEN if [ -f "${s3SecretKeyPathNode}" ]; then set +x @@ -5302,6 +5316,7 @@ def launchTestJobs(pipeline, testFilter) "DGX_B300-4_GPUs-PyTorch-Post-Merge-2": ["auto:dgx-b300-flex", "l0_dgx_b300", 2, 2, 4, 1, true], // VisualGen PerfSanity post-merge test "DGX_B200-8_GPUs-PyTorch-VisualGen-PerfSanity-Post-Merge-1": ["auto:dgx-b200-flex", "l0_b200_visual_gen_perf_sanity", 1, 1, 8, 1, true], + "DGX_B200-16_GPUs-2_Nodes-PyTorch-VisualGen-Post-Merge-1": ["auto:dgx-b200-flex", "l0_b200_visual_gen_multinode", 1, 1, 16, 2], // Single-GPU Gemma4 PerfSanity regression gate and baseline "DGX_B200-PyTorch-PerfSanity-1": ["auto:dgx-b200-flex", "l0_b200_perf_sanity", 1, 1, 1, 1, true], // PerfSanity post-merge tests diff --git a/jenkins/scripts/slurm_run.sh b/jenkins/scripts/slurm_run.sh index 25622daff391..e5791375349f 100755 --- a/jenkins/scripts/slurm_run.sh +++ b/jenkins/scripts/slurm_run.sh @@ -56,6 +56,26 @@ env | sort echo "Full Command: $pytestCommand" +if [[ "${TRTLLM_VISUAL_GEN_MULTINODE_SLURM_PARENT:-0}" == "1" ]]; then + install_done_dir="${jobWorkspace}/visual_gen_multinode_install_done" + mkdir -p "$install_done_dir" + if [[ "${SLURM_LOCALID:-0}" == "0" ]]; then + touch "${install_done_dir}/node_${SLURM_NODEID:-0}" + fi + + if [[ "${SLURM_PROCID:-0}" == "0" ]]; then + expected_nodes="${SLURM_JOB_NUM_NODES:-${SLURM_NNODES:-1}}" + while [[ "$(find "$install_done_dir" -maxdepth 1 -type f -name 'node_*' | wc -l)" -lt "$expected_nodes" ]]; do + echo "Waiting for VisualGen multi-node install markers in $install_done_dir" + sleep 10 + done + pytestCommand="env -u SLURM_PROCID -u SLURM_NTASKS -u SLURM_LOCALID -u SLURM_NODEID -u SLURM_GTIDS ${pytestCommand}" + else + echo "Rank${SLURM_PROCID} finished setup; rank0 will run the VisualGen SLURM parent pytest" + exit 0 + fi +fi + # For single-node test runs or disaggregated benchmark/server runs, clear all # environment variables related to Slurm and MPI. This prevents test processes # (e.g., pytest) from incorrectly initializing MPI when running under a diff --git a/tests/integration/defs/examples/visual_gen/test_visual_gen_multi_gpu.py b/tests/integration/defs/examples/visual_gen/test_visual_gen_multi_gpu.py index 0c30df8d2398..dd8e99477e67 100644 --- a/tests/integration/defs/examples/visual_gen/test_visual_gen_multi_gpu.py +++ b/tests/integration/defs/examples/visual_gen/test_visual_gen_multi_gpu.py @@ -340,8 +340,8 @@ def _run_wan22_multinode_slurm_parent(variant_name): os.environ.get("TLLM_SPAWN_PROXY_PROCESS") == "1" and _trtllm_launch_wrapper_world_size() >= WAN22_LPIPS_MULTINODE_WORLD_SIZE ): - pytest.skip( - "VisualGen SLURM external-launch coverage is skipped under " + pytest.fail( + "VisualGen SLURM external-launch coverage cannot run under " "trtllm-llmapi-launch because that wrapper removes SLURM_* env " "before user code. Run this nodeid with direct srun so " "_detect_external_launch() sees the real SLURM rank environment." diff --git a/tests/integration/test_lists/test-db/l0_b200_visual_gen_multinode.yml b/tests/integration/test_lists/test-db/l0_b200_visual_gen_multinode.yml new file mode 100644 index 000000000000..c8fb8caec701 --- /dev/null +++ b/tests/integration/test_lists/test-db/l0_b200_visual_gen_multinode.yml @@ -0,0 +1,33 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +version: 0.0.1 +l0_b200_visual_gen_multinode: +- condition: + ranges: + # 2 nodes with each node has 8 GPUs. + system_gpu_count: + gte: 16 + lte: 16 + wildcards: + gpu: + - '*b200*' + linux_distribution_name: ubuntu* + cpu: x86_64 + terms: + stage: post_merge + backend: pytorch + tests: + - examples/visual_gen/test_visual_gen_multi_gpu.py::test_wan22_t2v_lpips_against_golden_multinode_slurm[cfg2_attn2d_2x2_ulysses2] TIMEOUT (3600) From 714f085ad31fc9b784133e08b60e2563039c1d79 Mon Sep 17 00:00:00 2001 From: Ying Guo <244492186+yingguo-trt@users.noreply.github.com> Date: Tue, 21 Jul 2026 11:27:31 +0800 Subject: [PATCH 04/22] [TRTLLM-13304][test] Cover VisualGen external-launch init and rank-0 wiring MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The external-launch branch of VisualGen.__init__ (world-size validation, rank!=0 pure-worker exit) and the rank-0 DiffusionRemoteClient wiring (ZMQ connect addresses must use the launcher's master_addr, no local mp.Process spawn, own worker in a daemon thread with in_client_process=True) had no unit coverage — a wrong connect address keeps every single-node test green and only fails on a real 2-node run. Add CPU-only mocked tests following the existing patterns in test_visual_gen_multinode.py. Signed-off-by: Ying Guo <244492186+yingguo-trt@users.noreply.github.com> --- .../multi_gpu/test_visual_gen_multinode.py | 135 +++++++++++++++++- 1 file changed, 134 insertions(+), 1 deletion(-) diff --git a/tests/unittest/_torch/visual_gen/multi_gpu/test_visual_gen_multinode.py b/tests/unittest/_torch/visual_gen/multi_gpu/test_visual_gen_multinode.py index fddadb291d60..8e9804792a4e 100644 --- a/tests/unittest/_torch/visual_gen/multi_gpu/test_visual_gen_multinode.py +++ b/tests/unittest/_torch/visual_gen/multi_gpu/test_visual_gen_multinode.py @@ -13,7 +13,11 @@ from tensorrt_llm._torch.visual_gen.executor import run_diffusion_worker from tensorrt_llm.visual_gen.args import VisualGenArgs -from tensorrt_llm.visual_gen.visual_gen import DiffusionRemoteClient, _detect_external_launch +from tensorrt_llm.visual_gen.visual_gen import ( + DiffusionRemoteClient, + VisualGen, + _detect_external_launch, +) # ============================================================================= # _detect_external_launch() @@ -248,3 +252,132 @@ def pre_set_event(): f"Worker {i}: expected local_rank={i}, got {kwargs['local_rank']}. " f"With LOCAL_RANK=0 in env, all workers would get device_id=0." ) + + +# ============================================================================= +# VisualGen.__init__ — external launch branch +# ============================================================================= + + +class TestVisualGenExternalLaunchInit: + """VisualGen.__init__ external-launch handling — no GPU required. + + Covers the branch where torchrun/srun launched the process: world-size + validation and the rank != 0 pure-worker path. Neither is reachable from + single-node tests, and the 2-node E2E only runs post-merge. + """ + + def test_world_size_mismatch_raises(self): + """Launcher task count must match parallel_config.n_workers.""" + args = VisualGenArgs(model="/tmp/model", parallel_config={"ulysses_size": 4}) + + with patch( + "tensorrt_llm.visual_gen.visual_gen._detect_external_launch", + return_value=(0, 0, 8, "node0", 29500), + ): + with pytest.raises(ValueError, match="does not match"): + VisualGen(model="/tmp/model", args=args) + + def test_nonzero_rank_runs_worker_then_exits(self): + """Rank != 0 must run the worker loop and exit — never reach user code.""" + args = VisualGenArgs(model="/tmp/model", parallel_config={"ulysses_size": 4}) + mock_worker = MagicMock() + + with ( + patch( + "tensorrt_llm.visual_gen.visual_gen._detect_external_launch", + return_value=(2, 2, 4, "node0", 29500), + ), + patch("tensorrt_llm.visual_gen.visual_gen.run_diffusion_worker", mock_worker), + ): + with pytest.raises(SystemExit) as exc_info: + VisualGen(model="/tmp/model", args=args) + + assert exc_info.value.code == 0 + mock_worker.assert_called_once() + kwargs = mock_worker.call_args.kwargs + assert kwargs["rank"] == 2 + assert kwargs["local_rank"] == 2 + assert kwargs["world_size"] == 4 + assert kwargs["master_addr"] == "node0" + assert kwargs["master_port"] == 29500 + # Non-zero ranks receive requests via dist broadcast, not ZMQ + assert kwargs["request_queue_addr"] is None + assert kwargs["response_queue_addr"] is None + + +# ============================================================================= +# DiffusionRemoteClient — external launch wiring (rank 0 coordinator) +# ============================================================================= + + +class TestExternalLaunchClientWiring: + """Rank-0 coordinator wiring in external-launch mode — no GPU required. + + Regression target: the ZMQ connect addresses handed to workers must use + the launcher's master_addr, not 127.0.0.1 / local host IP. Getting this + wrong keeps every single-node test green while remote-node workers can + never connect — only a real 2-node run would catch it. + """ + + @pytest.fixture(autouse=True) + def _clean_env(self, monkeypatch): + for var in ["RANK", "WORLD_SIZE", "LOCAL_RANK", "SLURM_PROCID", "SLURM_NTASKS"]: + monkeypatch.delenv(var, raising=False) + + def test_rank0_uses_master_addr_and_thread_worker(self): + args = VisualGenArgs(model="/tmp/model", parallel_config={"ulysses_size": 4}) + + mock_ctx = MagicMock() + original_event = threading.Event + + def pre_set_event(): + e = original_event() + e.set() + return e + + with ( + patch( + "tensorrt_llm._torch.visual_gen.executor._detect_external_launch", + return_value=(0, 0, 4, "node0", 29500), + ), + patch("tensorrt_llm._torch.visual_gen.executor.mp.get_context", return_value=mock_ctx), + patch("tensorrt_llm._torch.visual_gen.executor.threading.Thread") as mock_thread_cls, + patch( + "tensorrt_llm._torch.visual_gen.executor.threading.Event", side_effect=pre_set_event + ), + patch.object(DiffusionRemoteClient, "_wait_ready"), + ): + mock_thread_cls.return_value = MagicMock() + client = DiffusionRemoteClient(args=args) + + # Distributed rendezvous comes from the launcher, not a local port + assert client.master_addr == "node0" + assert client.master_port == 29500 + + # Remote-node workers connect to rank 0 via master_addr + assert client.req_addr_connect.startswith("tcp://node0:") + assert client.resp_addr_connect.startswith("tcp://node0:") + + # No local mp.Process spawn — other ranks were launched externally + mock_ctx.Process.assert_not_called() + assert client.worker_processes == [] + + # Rank 0's own worker runs in a daemon thread inside the client process + worker_calls = [ + c + for c in mock_thread_cls.call_args_list + if "in_client_process" in c.kwargs.get("kwargs", {}) + ] + assert len(worker_calls) == 1, "Expected exactly one external-launch worker thread" + thread_call = worker_calls[0] + assert thread_call.kwargs["daemon"] is True + worker_kwargs = thread_call.kwargs["kwargs"] + assert worker_kwargs["in_client_process"] is True + assert worker_kwargs["rank"] == 0 + assert worker_kwargs["local_rank"] == 0 + assert worker_kwargs["world_size"] == 4 + assert worker_kwargs["master_addr"] == "node0" + assert worker_kwargs["master_port"] == 29500 + assert worker_kwargs["request_queue_addr"] == client.req_addr_connect + assert worker_kwargs["response_queue_addr"] == client.resp_addr_connect From 5fd4dce26125ca323ab4666fafa3849a09a1deab Mon Sep 17 00:00:00 2001 From: Ying Guo <244492186+yingguo-trt@users.noreply.github.com> Date: Tue, 21 Jul 2026 11:46:23 +0800 Subject: [PATCH 05/22] [TRTLLM-13304][fix] Fix NameError in VisualGen multinode SLURM rank path ParallelConfig is only imported lazily inside _parallel_config() (kept lazy for mp.spawn sys.path ordering), so the bare ParallelConfig reference in _run_wan22_multinode_slurm_rank raised NameError on every SLURM rank before VisualGen was even constructed. Use the _parallel_config helper like the other call sites. Signed-off-by: Ying Guo <244492186+yingguo-trt@users.noreply.github.com> --- .../defs/examples/visual_gen/test_visual_gen_multi_gpu.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/integration/defs/examples/visual_gen/test_visual_gen_multi_gpu.py b/tests/integration/defs/examples/visual_gen/test_visual_gen_multi_gpu.py index dd8e99477e67..bfb195d721f4 100644 --- a/tests/integration/defs/examples/visual_gen/test_visual_gen_multi_gpu.py +++ b/tests/integration/defs/examples/visual_gen/test_visual_gen_multi_gpu.py @@ -273,7 +273,7 @@ def _run_wan22_multinode_slurm_rank(tmp_path, variant_name, parallel): ) _ensure_slurm_external_launch_env() - ParallelConfig(**parallel).validate_world_size(world_size) + _parallel_config(**parallel).validate_world_size(world_size) model_path = _lpips_model_path("Wan2.2-T2V-A14B-Diffusers") _skip_if_missing(model_path, "Wan 2.2 checkpoint", is_dir=True) From d611f879f6043fdd0a32027522c9249aed7086a8 Mon Sep 17 00:00:00 2001 From: Ying Guo <244492186+yingguo-trt@users.noreply.github.com> Date: Tue, 21 Jul 2026 11:57:36 +0800 Subject: [PATCH 06/22] [TRTLLM-13304][fix] Fix undefined MODULES_AVAILABLE and missing media deps in multinode case Two blockers in the VisualGen multinode SLURM case: - MODULES_AVAILABLE was referenced at the top of the SLURM rank path but never defined in this file (ruff F821), so every rank raised NameError before constructing VisualGen. Drop the stale guard copied from the unittest harness; module-level imports already fail loudly. - The multinode nodeid never requested _visual_gen_deps, so a clean node lacked av/ffmpeg for rank 0's MP4 save and video LPIPS eval. Resolve the fixture parent-side only via request.getfixturevalue: rank 0 shares node 0's venv with the parent, remote workers never touch media deps, and the bare child pytest run has no llm_venv harness (and 16 ranks racing pip/apt would be unsafe). Signed-off-by: Ying Guo <244492186+yingguo-trt@users.noreply.github.com> --- .../examples/visual_gen/test_visual_gen_multi_gpu.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/tests/integration/defs/examples/visual_gen/test_visual_gen_multi_gpu.py b/tests/integration/defs/examples/visual_gen/test_visual_gen_multi_gpu.py index bfb195d721f4..2e8633c0225e 100644 --- a/tests/integration/defs/examples/visual_gen/test_visual_gen_multi_gpu.py +++ b/tests/integration/defs/examples/visual_gen/test_visual_gen_multi_gpu.py @@ -257,9 +257,6 @@ def _ensure_slurm_external_launch_env(): def _run_wan22_multinode_slurm_rank(tmp_path, variant_name, parallel): - if not MODULES_AVAILABLE: - pytest.skip("Required modules not available") - rank_env = _slurm_rank_env() if rank_env is None: pytest.skip("This VisualGen multi-node case must run under SLURM rank env") @@ -524,7 +521,7 @@ def test_wan22_t2v_lpips_against_golden_tp(_visual_gen_deps, tmp_path, variant_n WAN22_LPIPS_MULTINODE_VARIANTS, ids=[name for name, _ in WAN22_LPIPS_MULTINODE_VARIANTS], ) -def test_wan22_t2v_lpips_against_golden_multinode_slurm(tmp_path, variant_name, parallel): +def test_wan22_t2v_lpips_against_golden_multinode_slurm(request, tmp_path, variant_name, parallel): if _slurm_rank_env() is not None: _run_wan22_multinode_slurm_rank(tmp_path, variant_name, parallel) return @@ -532,4 +529,9 @@ def test_wan22_t2v_lpips_against_golden_multinode_slurm(tmp_path, variant_name, if os.environ.get(_MULTINODE_SLURM_CHILD_ENV): pytest.skip("VisualGen SLURM child was not launched with SLURM rank env") + # Media deps (av / ffmpeg) are installed parent-side only: rank 0 shares + # node 0's venv with the parent, remote-node workers never touch media, + # and the bare child pytest invocation has no llm_venv harness to run + # the fixture (16 ranks racing pip/apt would be unsafe anyway). + request.getfixturevalue("_visual_gen_deps") _run_wan22_multinode_slurm_parent(variant_name) From a1e637426a5d78187059b3e49c8f51f57b23eccf Mon Sep 17 00:00:00 2001 From: Ying Guo <244492186+yingguo-trt@users.noreply.github.com> Date: Tue, 21 Jul 2026 12:04:53 +0800 Subject: [PATCH 07/22] [TRTLLM-13304][fix] Fix multinode case timeout layering MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Test-list TIMEOUT is in minutes (test_list_parser.py multiplies by 60), so TIMEOUT (3600) on the multinode nodeid meant 60 hours — a hung run would pin a 2-node 16xB200 allocation for days. Set it to 90 minutes. The QA multinode list entry carries no TIMEOUT marker and falls back to the pipeline's outer --timeout=3600s, which would fire before the inner srun subprocess timeout (same 3600s) and lose its captured output. Lower the inner default to 3300s so the srun timeout with output always fires first. Signed-off-by: Ying Guo <244492186+yingguo-trt@users.noreply.github.com> --- .../defs/examples/visual_gen/test_visual_gen_multi_gpu.py | 4 +++- .../test_lists/test-db/l0_b200_visual_gen_multinode.yml | 2 +- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/tests/integration/defs/examples/visual_gen/test_visual_gen_multi_gpu.py b/tests/integration/defs/examples/visual_gen/test_visual_gen_multi_gpu.py index 2e8633c0225e..52d998f5fae0 100644 --- a/tests/integration/defs/examples/visual_gen/test_visual_gen_multi_gpu.py +++ b/tests/integration/defs/examples/visual_gen/test_visual_gen_multi_gpu.py @@ -216,7 +216,9 @@ def _trtllm_launch_wrapper_world_size(): def _multinode_subprocess_timeout(): - return int(os.environ.get("TRTLLM_VISUAL_GEN_MULTINODE_TIMEOUT", "3600")) + # Default below the QA pipeline's outer --timeout=3600s so the inner srun + # timeout (with captured output) fires before pytest-timeout kills the test. + return int(os.environ.get("TRTLLM_VISUAL_GEN_MULTINODE_TIMEOUT", "3300")) def _resolve_slurm_master_addr(): diff --git a/tests/integration/test_lists/test-db/l0_b200_visual_gen_multinode.yml b/tests/integration/test_lists/test-db/l0_b200_visual_gen_multinode.yml index c8fb8caec701..880fd9380284 100644 --- a/tests/integration/test_lists/test-db/l0_b200_visual_gen_multinode.yml +++ b/tests/integration/test_lists/test-db/l0_b200_visual_gen_multinode.yml @@ -30,4 +30,4 @@ l0_b200_visual_gen_multinode: stage: post_merge backend: pytorch tests: - - examples/visual_gen/test_visual_gen_multi_gpu.py::test_wan22_t2v_lpips_against_golden_multinode_slurm[cfg2_attn2d_2x2_ulysses2] TIMEOUT (3600) + - examples/visual_gen/test_visual_gen_multi_gpu.py::test_wan22_t2v_lpips_against_golden_multinode_slurm[cfg2_attn2d_2x2_ulysses2] TIMEOUT (90) From 5e567546746f9b679d58517b09733a060b876ae8 Mon Sep 17 00:00:00 2001 From: Ying Guo <244492186+yingguo-trt@users.noreply.github.com> Date: Tue, 21 Jul 2026 12:18:50 +0800 Subject: [PATCH 08/22] [TRTLLM-13304][test] Simplify multinode case to single variant and tidy unit tests MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The multinode SLURM case has exactly one parallel variant, so drop the parametrization in favor of a WAN22_LPIPS_MULTINODE_PARALLEL constant and sync the nodeid in the dedicated L0 stage. Remove the QA multinode list entry: that pipeline wraps pytest with trtllm-llmapi-launch, which strips SLURM_* env before user code, so the case would always hard-fail there — the dedicated post-merge SLURM stage remains the single consumer. Compress assertions and drop banner comments in the multinode unit tests. Signed-off-by: Ying Guo <244492186+yingguo-trt@users.noreply.github.com> --- .../visual_gen/test_visual_gen_multi_gpu.py | 36 ++++----- .../test_lists/qa/llm_function_multinode.txt | 1 - .../test-db/l0_b200_visual_gen_multinode.yml | 2 +- .../multi_gpu/test_visual_gen_multinode.py | 77 +++++-------------- 4 files changed, 37 insertions(+), 79 deletions(-) diff --git a/tests/integration/defs/examples/visual_gen/test_visual_gen_multi_gpu.py b/tests/integration/defs/examples/visual_gen/test_visual_gen_multi_gpu.py index 52d998f5fae0..148e6245a603 100644 --- a/tests/integration/defs/examples/visual_gen/test_visual_gen_multi_gpu.py +++ b/tests/integration/defs/examples/visual_gen/test_visual_gen_multi_gpu.py @@ -79,12 +79,11 @@ def _parallel_config(**kwargs): WAN22_LPIPS_MULTINODE_WORLD_SIZE = 16 WAN22_LPIPS_MULTINODE_NODES = 2 WAN22_LPIPS_MULTINODE_GPUS_PER_NODE = 8 -WAN22_LPIPS_MULTINODE_VARIANTS = [ - ( - "cfg2_attn2d_2x2_ulysses2", - {"cfg_size": 2, "attn2d_size": (2, 2), "ulysses_size": 2}, - ), -] +WAN22_LPIPS_MULTINODE_PARALLEL = { + "cfg_size": 2, + "attn2d_size": (2, 2), + "ulysses_size": 2, +} _MULTINODE_SLURM_CHILD_ENV = "TRTLLM_VISUAL_GEN_MULTINODE_SLURM_CHILD" @@ -258,7 +257,7 @@ def _ensure_slurm_external_launch_env(): os.environ.pop(var, None) -def _run_wan22_multinode_slurm_rank(tmp_path, variant_name, parallel): +def _run_wan22_multinode_slurm_rank(tmp_path): rank_env = _slurm_rank_env() if rank_env is None: pytest.skip("This VisualGen multi-node case must run under SLURM rank env") @@ -272,7 +271,7 @@ def _run_wan22_multinode_slurm_rank(tmp_path, variant_name, parallel): ) _ensure_slurm_external_launch_env() - _parallel_config(**parallel).validate_world_size(world_size) + _parallel_config(**WAN22_LPIPS_MULTINODE_PARALLEL).validate_world_size(world_size) model_path = _lpips_model_path("Wan2.2-T2V-A14B-Diffusers") _skip_if_missing(model_path, "Wan 2.2 checkpoint", is_dir=True) @@ -284,7 +283,7 @@ def _run_wan22_multinode_slurm_rank(tmp_path, variant_name, parallel): compilation_config=CompilationConfig(skip_warmup=True), torch_compile_config=TorchCompileConfig(enable=False), attention_config=AttentionConfig(backend="FA4"), - parallel_config=parallel, + parallel_config=WAN22_LPIPS_MULTINODE_PARALLEL, ) visual_gen = None @@ -311,7 +310,7 @@ def _run_wan22_multinode_slurm_rank(tmp_path, variant_name, parallel): assert output.error is None, f"unexpected error on Wan 2.2 multi-node run: {output.error}" assert output.video is not None - generated_path = tmp_path / f"wan22_t2v_generated_{variant_name}_slurm.mp4" + generated_path = tmp_path / "wan22_t2v_generated_multinode_slurm.mp4" output.save(generated_path, frame_rate=WAN22_LPIPS_FRAME_RATE) assert generated_path.is_file(), ( f"VisualGen multi-node run did not produce {generated_path}" @@ -322,7 +321,7 @@ def _run_wan22_multinode_slurm_rank(tmp_path, variant_name, parallel): ) score = _run_lpips_eval( tmp_path, - f"wan22_t2v_{variant_name}_slurm", + "wan22_t2v_multinode_slurm", "video", WAN22_LPIPS_PROMPT, golden_path, @@ -334,7 +333,7 @@ def _run_wan22_multinode_slurm_rank(tmp_path, variant_name, parallel): visual_gen.shutdown() -def _run_wan22_multinode_slurm_parent(variant_name): +def _run_wan22_multinode_slurm_parent(): if ( os.environ.get("TLLM_SPAWN_PROXY_PROCESS") == "1" and _trtllm_launch_wrapper_world_size() >= WAN22_LPIPS_MULTINODE_WORLD_SIZE @@ -374,7 +373,7 @@ def _run_wan22_multinode_slurm_parent(variant_name): env.pop(var, None) test_file = str(Path(__file__).resolve()) - nodeid = f"{test_file}::test_wan22_t2v_lpips_against_golden_multinode_slurm[{variant_name}]" + nodeid = f"{test_file}::test_wan22_t2v_lpips_against_golden_multinode_slurm" cmd = [ "srun", "-l", @@ -518,14 +517,9 @@ def test_wan22_t2v_lpips_against_golden_tp(_visual_gen_deps, tmp_path, variant_n _run_wan22_t2v_lpips_case(tmp_path, variant_name, parallel) -@pytest.mark.parametrize( - "variant_name,parallel", - WAN22_LPIPS_MULTINODE_VARIANTS, - ids=[name for name, _ in WAN22_LPIPS_MULTINODE_VARIANTS], -) -def test_wan22_t2v_lpips_against_golden_multinode_slurm(request, tmp_path, variant_name, parallel): +def test_wan22_t2v_lpips_against_golden_multinode_slurm(request, tmp_path): if _slurm_rank_env() is not None: - _run_wan22_multinode_slurm_rank(tmp_path, variant_name, parallel) + _run_wan22_multinode_slurm_rank(tmp_path) return if os.environ.get(_MULTINODE_SLURM_CHILD_ENV): @@ -536,4 +530,4 @@ def test_wan22_t2v_lpips_against_golden_multinode_slurm(request, tmp_path, varia # and the bare child pytest invocation has no llm_venv harness to run # the fixture (16 ranks racing pip/apt would be unsafe anyway). request.getfixturevalue("_visual_gen_deps") - _run_wan22_multinode_slurm_parent(variant_name) + _run_wan22_multinode_slurm_parent() diff --git a/tests/integration/test_lists/qa/llm_function_multinode.txt b/tests/integration/test_lists/qa/llm_function_multinode.txt index 10ca417883d9..cb10040d1934 100644 --- a/tests/integration/test_lists/qa/llm_function_multinode.txt +++ b/tests/integration/test_lists/qa/llm_function_multinode.txt @@ -7,4 +7,3 @@ test_e2e.py::test_multi_nodes_eval[MiniMax-M3-tp16-mmlu] test_e2e.py::test_openai_disagg_multi_nodes_completion[ctx_tp2pp1-gen_tp2pp1] test_e2e.py::test_openai_disagg_multi_nodes_completion[ctx_tp1pp2-gen_tp1pp2] test_e2e.py::test_openai_disagg_multi_nodes_completion_service_discovery[etcd] -examples/visual_gen/test_visual_gen_multi_gpu.py::test_wan22_t2v_lpips_against_golden_multinode_slurm[cfg2_attn2d_2x2_ulysses2] diff --git a/tests/integration/test_lists/test-db/l0_b200_visual_gen_multinode.yml b/tests/integration/test_lists/test-db/l0_b200_visual_gen_multinode.yml index 880fd9380284..f3cf16d02bb1 100644 --- a/tests/integration/test_lists/test-db/l0_b200_visual_gen_multinode.yml +++ b/tests/integration/test_lists/test-db/l0_b200_visual_gen_multinode.yml @@ -30,4 +30,4 @@ l0_b200_visual_gen_multinode: stage: post_merge backend: pytorch tests: - - examples/visual_gen/test_visual_gen_multi_gpu.py::test_wan22_t2v_lpips_against_golden_multinode_slurm[cfg2_attn2d_2x2_ulysses2] TIMEOUT (90) + - examples/visual_gen/test_visual_gen_multi_gpu.py::test_wan22_t2v_lpips_against_golden_multinode_slurm TIMEOUT (90) diff --git a/tests/unittest/_torch/visual_gen/multi_gpu/test_visual_gen_multinode.py b/tests/unittest/_torch/visual_gen/multi_gpu/test_visual_gen_multinode.py index 8e9804792a4e..3bdfdec1cffc 100644 --- a/tests/unittest/_torch/visual_gen/multi_gpu/test_visual_gen_multinode.py +++ b/tests/unittest/_torch/visual_gen/multi_gpu/test_visual_gen_multinode.py @@ -254,21 +254,8 @@ def pre_set_event(): ) -# ============================================================================= -# VisualGen.__init__ — external launch branch -# ============================================================================= - - class TestVisualGenExternalLaunchInit: - """VisualGen.__init__ external-launch handling — no GPU required. - - Covers the branch where torchrun/srun launched the process: world-size - validation and the rank != 0 pure-worker path. Neither is reachable from - single-node tests, and the 2-node E2E only runs post-merge. - """ - def test_world_size_mismatch_raises(self): - """Launcher task count must match parallel_config.n_workers.""" args = VisualGenArgs(model="/tmp/model", parallel_config={"ulysses_size": 4}) with patch( @@ -279,7 +266,6 @@ def test_world_size_mismatch_raises(self): VisualGen(model="/tmp/model", args=args) def test_nonzero_rank_runs_worker_then_exits(self): - """Rank != 0 must run the worker loop and exit — never reach user code.""" args = VisualGenArgs(model="/tmp/model", parallel_config={"ulysses_size": 4}) mock_worker = MagicMock() @@ -294,32 +280,18 @@ def test_nonzero_rank_runs_worker_then_exits(self): VisualGen(model="/tmp/model", args=args) assert exc_info.value.code == 0 - mock_worker.assert_called_once() kwargs = mock_worker.call_args.kwargs - assert kwargs["rank"] == 2 - assert kwargs["local_rank"] == 2 - assert kwargs["world_size"] == 4 - assert kwargs["master_addr"] == "node0" - assert kwargs["master_port"] == 29500 - # Non-zero ranks receive requests via dist broadcast, not ZMQ - assert kwargs["request_queue_addr"] is None - assert kwargs["response_queue_addr"] is None - - -# ============================================================================= -# DiffusionRemoteClient — external launch wiring (rank 0 coordinator) -# ============================================================================= + assert ( + kwargs["rank"], + kwargs["local_rank"], + kwargs["world_size"], + kwargs["master_addr"], + kwargs["master_port"], + ) == (2, 2, 4, "node0", 29500) + assert kwargs["request_queue_addr"] is kwargs["response_queue_addr"] is None class TestExternalLaunchClientWiring: - """Rank-0 coordinator wiring in external-launch mode — no GPU required. - - Regression target: the ZMQ connect addresses handed to workers must use - the launcher's master_addr, not 127.0.0.1 / local host IP. Getting this - wrong keeps every single-node test green while remote-node workers can - never connect — only a real 2-node run would catch it. - """ - @pytest.fixture(autouse=True) def _clean_env(self, monkeypatch): for var in ["RANK", "WORLD_SIZE", "LOCAL_RANK", "SLURM_PROCID", "SLURM_NTASKS"]: @@ -332,9 +304,9 @@ def test_rank0_uses_master_addr_and_thread_worker(self): original_event = threading.Event def pre_set_event(): - e = original_event() - e.set() - return e + event = original_event() + event.set() + return event with ( patch( @@ -351,33 +323,26 @@ def pre_set_event(): mock_thread_cls.return_value = MagicMock() client = DiffusionRemoteClient(args=args) - # Distributed rendezvous comes from the launcher, not a local port - assert client.master_addr == "node0" - assert client.master_port == 29500 - - # Remote-node workers connect to rank 0 via master_addr + assert (client.master_addr, client.master_port) == ("node0", 29500) assert client.req_addr_connect.startswith("tcp://node0:") assert client.resp_addr_connect.startswith("tcp://node0:") - - # No local mp.Process spawn — other ranks were launched externally mock_ctx.Process.assert_not_called() - assert client.worker_processes == [] - - # Rank 0's own worker runs in a daemon thread inside the client process worker_calls = [ c for c in mock_thread_cls.call_args_list if "in_client_process" in c.kwargs.get("kwargs", {}) ] - assert len(worker_calls) == 1, "Expected exactly one external-launch worker thread" + assert len(worker_calls) == 1 thread_call = worker_calls[0] assert thread_call.kwargs["daemon"] is True worker_kwargs = thread_call.kwargs["kwargs"] - assert worker_kwargs["in_client_process"] is True - assert worker_kwargs["rank"] == 0 - assert worker_kwargs["local_rank"] == 0 - assert worker_kwargs["world_size"] == 4 - assert worker_kwargs["master_addr"] == "node0" - assert worker_kwargs["master_port"] == 29500 + assert worker_kwargs["in_client_process"] + assert ( + worker_kwargs["rank"], + worker_kwargs["local_rank"], + worker_kwargs["world_size"], + worker_kwargs["master_addr"], + worker_kwargs["master_port"], + ) == (0, 0, 4, "node0", 29500) assert worker_kwargs["request_queue_addr"] == client.req_addr_connect assert worker_kwargs["response_queue_addr"] == client.resp_addr_connect From 8c9c2332dfaaa6c541c7f7d14c6c664ec1fa460b Mon Sep 17 00:00:00 2001 From: Ying Guo <244492186+yingguo-trt@users.noreply.github.com> Date: Tue, 21 Jul 2026 12:35:39 +0800 Subject: [PATCH 09/22] [TRTLLM-13304][test] Re-register VisualGen multinode case in QA multinode list The QA cluster pipeline (llm_test_cluster) runs this list with direct srun (-N2 --ntasks-per-node=8 = 16 ranks) and already opts multinode_slurm cases out of the trtllm-llmapi-launch wrapper, so the SLURM rank environment reaches the test and every rank takes the external-launch rank path. The earlier removal assumed the wrapper always applies; re-add the entry with the de-parametrized nodeid. Signed-off-by: Ying Guo <244492186+yingguo-trt@users.noreply.github.com> --- tests/integration/test_lists/qa/llm_function_multinode.txt | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/integration/test_lists/qa/llm_function_multinode.txt b/tests/integration/test_lists/qa/llm_function_multinode.txt index cb10040d1934..271f5d295432 100644 --- a/tests/integration/test_lists/qa/llm_function_multinode.txt +++ b/tests/integration/test_lists/qa/llm_function_multinode.txt @@ -7,3 +7,4 @@ test_e2e.py::test_multi_nodes_eval[MiniMax-M3-tp16-mmlu] test_e2e.py::test_openai_disagg_multi_nodes_completion[ctx_tp2pp1-gen_tp2pp1] test_e2e.py::test_openai_disagg_multi_nodes_completion[ctx_tp1pp2-gen_tp1pp2] test_e2e.py::test_openai_disagg_multi_nodes_completion_service_discovery[etcd] +examples/visual_gen/test_visual_gen_multi_gpu.py::test_wan22_t2v_lpips_against_golden_multinode_slurm From ea4a80a3ade2e181676d505d29658ee6b0c90d0e Mon Sep 17 00:00:00 2001 From: Ying Guo <244492186+yingguo-trt@users.noreply.github.com> Date: Tue, 21 Jul 2026 17:42:10 +0800 Subject: [PATCH 10/22] [TRTLLM-13304][fix] Clean up leaked threads in VisualGen multinode case pytest-threadleak flags three daemon threads after the multinode case: InductorSubproc / subproc_worker_timer (unconditional @torch.compile spawns the inductor worker pool; same root cause as nvbugs/6215688) and rank 0's run_diffusion_worker in-client thread (shutdown()'s bounded join can be outlasted by 16-rank dist teardown). Reuse the existing fix pattern from test_visual_gen.py: disable the inductor quiesce timer up front and call _cleanup_cuda() (shutdown_compile_workers) in the shared finally so every rank cleans up, and explicitly join rank 0's worker thread after shutdown. Signed-off-by: Ying Guo <244492186+yingguo-trt@users.noreply.github.com> --- .../examples/visual_gen/test_visual_gen_multi_gpu.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/tests/integration/defs/examples/visual_gen/test_visual_gen_multi_gpu.py b/tests/integration/defs/examples/visual_gen/test_visual_gen_multi_gpu.py index 148e6245a603..f3bab9b37a54 100644 --- a/tests/integration/defs/examples/visual_gen/test_visual_gen_multi_gpu.py +++ b/tests/integration/defs/examples/visual_gen/test_visual_gen_multi_gpu.py @@ -37,6 +37,8 @@ WAN22_LPIPS_SEED, WAN22_LPIPS_WIDTH, _assert_lpips_below_threshold, + _cleanup_cuda, + _disable_inductor_compile_worker_quiesce, _golden_media_path, _lpips_model_path, _run_lpips_eval, @@ -271,6 +273,7 @@ def _run_wan22_multinode_slurm_rank(tmp_path): ) _ensure_slurm_external_launch_env() + _disable_inductor_compile_worker_quiesce() _parallel_config(**WAN22_LPIPS_MULTINODE_PARALLEL).validate_world_size(world_size) model_path = _lpips_model_path("Wan2.2-T2V-A14B-Diffusers") _skip_if_missing(model_path, "Wan 2.2 checkpoint", is_dir=True) @@ -331,6 +334,13 @@ def _run_wan22_multinode_slurm_rank(tmp_path): finally: if visual_gen is not None: visual_gen.shutdown() + # shutdown() joins rank 0's in-client worker thread with a bounded + # timeout; 16-rank dist teardown can outlast it and trip + # pytest-threadleak, so wait for the thread explicitly. + worker_thread = getattr(visual_gen.executor, "_ext_worker_thread", None) + if worker_thread is not None and worker_thread.is_alive(): + worker_thread.join(timeout=120) + _cleanup_cuda() def _run_wan22_multinode_slurm_parent(): From 80020164b643e082576dc1cd16306af0c94d4d79 Mon Sep 17 00:00:00 2001 From: Ying Guo <244492186+yingguo-trt@users.noreply.github.com> Date: Mon, 3 Aug 2026 14:52:19 +0800 Subject: [PATCH 11/22] [TRTLLM-13304][fix] Align multinode LPIPS golden with the multi-GPU baseline The multinode case runs the FA4 backend fully-eager like the single-node multi-GPU cases, but compared its output against the default-path golden, so the LPIPS score was not comparable. Use WAN22_MULTI_GPU_LPIPS_GOLDEN_VIDEO and wrap the run in _lpips_deterministic_algorithms(fully_eager=True) to match _wan22_lpips_distributed_worker, which the multinode path bypasses by calling VisualGen directly. Rename the test for consistency with the other LPIPS nodeids and drop the QA multinode list entry: the dedicated L0 post-merge stage already covers this case. Signed-off-by: Ying Guo <244492186+yingguo-trt@users.noreply.github.com> --- .../visual_gen/test_visual_gen_multi_gpu.py | 90 ++++++++++--------- .../test_lists/qa/llm_function_multinode.txt | 1 - .../test-db/l0_b200_visual_gen_multinode.yml | 2 +- 3 files changed, 49 insertions(+), 44 deletions(-) diff --git a/tests/integration/defs/examples/visual_gen/test_visual_gen_multi_gpu.py b/tests/integration/defs/examples/visual_gen/test_visual_gen_multi_gpu.py index f3bab9b37a54..c1f69adbe6ae 100644 --- a/tests/integration/defs/examples/visual_gen/test_visual_gen_multi_gpu.py +++ b/tests/integration/defs/examples/visual_gen/test_visual_gen_multi_gpu.py @@ -40,6 +40,7 @@ _cleanup_cuda, _disable_inductor_compile_worker_quiesce, _golden_media_path, + _lpips_deterministic_algorithms, _lpips_model_path, _run_lpips_eval, _run_wan_lpips_pipeline, @@ -291,46 +292,51 @@ def _run_wan22_multinode_slurm_rank(tmp_path): visual_gen = None try: - try: - visual_gen = VisualGen(model=model_path, args=visual_gen_args) - except SystemExit as exc: - assert rank != 0, "Only non-zero SLURM ranks should exit through worker mode" - assert exc.code in (0, None) - return - - assert rank == 0 - params = VisualGenParams( - height=WAN22_LPIPS_HEIGHT, - width=WAN22_LPIPS_WIDTH, - num_frames=WAN22_LPIPS_NUM_FRAMES, - num_inference_steps=WAN22_LPIPS_NUM_INFERENCE_STEPS, - guidance_scale=WAN22_LPIPS_GUIDANCE_SCALE, - seed=WAN22_LPIPS_SEED, - frame_rate=WAN22_LPIPS_FRAME_RATE, - negative_prompt=WAN22_LPIPS_NEGATIVE_PROMPT, - ) - output = visual_gen.generate(inputs=WAN22_LPIPS_PROMPT, params=params) - assert output.error is None, f"unexpected error on Wan 2.2 multi-node run: {output.error}" - assert output.video is not None - - generated_path = tmp_path / "wan22_t2v_generated_multinode_slurm.mp4" - output.save(generated_path, frame_rate=WAN22_LPIPS_FRAME_RATE) - assert generated_path.is_file(), ( - f"VisualGen multi-node run did not produce {generated_path}" - ) - - golden_path = _golden_media_path( - tmp_path, "wan22_t2v_lpips_golden_video.mp4", "Wan 2.2 LPIPS golden video" - ) - score = _run_lpips_eval( - tmp_path, - "wan22_t2v_multinode_slurm", - "video", - WAN22_LPIPS_PROMPT, - golden_path, - generated_path, - ) - _assert_lpips_below_threshold(score, WAN_MULTI_GPU_LPIPS_THRESHOLD) + with _lpips_deterministic_algorithms(fully_eager=True): + try: + visual_gen = VisualGen(model=model_path, args=visual_gen_args) + except SystemExit as exc: + assert rank != 0, "Only non-zero SLURM ranks should exit through worker mode" + assert exc.code in (0, None) + return + + assert rank == 0 + params = VisualGenParams( + height=WAN22_LPIPS_HEIGHT, + width=WAN22_LPIPS_WIDTH, + num_frames=WAN22_LPIPS_NUM_FRAMES, + num_inference_steps=WAN22_LPIPS_NUM_INFERENCE_STEPS, + guidance_scale=WAN22_LPIPS_GUIDANCE_SCALE, + seed=WAN22_LPIPS_SEED, + frame_rate=WAN22_LPIPS_FRAME_RATE, + negative_prompt=WAN22_LPIPS_NEGATIVE_PROMPT, + ) + output = visual_gen.generate(inputs=WAN22_LPIPS_PROMPT, params=params) + assert output.error is None, ( + f"unexpected error on Wan 2.2 multi-node run: {output.error}" + ) + assert output.video is not None + + generated_path = tmp_path / "wan22_t2v_generated_multinode_slurm.mp4" + output.save(generated_path, frame_rate=WAN22_LPIPS_FRAME_RATE) + assert generated_path.is_file(), ( + f"VisualGen multi-node run did not produce {generated_path}" + ) + + golden_path = _golden_media_path( + tmp_path, + WAN22_MULTI_GPU_LPIPS_GOLDEN_VIDEO, + "Wan 2.2 FA4 fully-eager LPIPS golden video", + ) + score = _run_lpips_eval( + tmp_path, + "wan22_t2v_multinode_slurm", + "video", + WAN22_LPIPS_PROMPT, + golden_path, + generated_path, + ) + _assert_lpips_below_threshold(score, WAN_MULTI_GPU_LPIPS_THRESHOLD) finally: if visual_gen is not None: visual_gen.shutdown() @@ -383,7 +389,7 @@ def _run_wan22_multinode_slurm_parent(): env.pop(var, None) test_file = str(Path(__file__).resolve()) - nodeid = f"{test_file}::test_wan22_t2v_lpips_against_golden_multinode_slurm" + nodeid = f"{test_file}::test_wan22_t2v_multinode_slurm_lpips_against_golden" cmd = [ "srun", "-l", @@ -527,7 +533,7 @@ def test_wan22_t2v_lpips_against_golden_tp(_visual_gen_deps, tmp_path, variant_n _run_wan22_t2v_lpips_case(tmp_path, variant_name, parallel) -def test_wan22_t2v_lpips_against_golden_multinode_slurm(request, tmp_path): +def test_wan22_t2v_multinode_slurm_lpips_against_golden(request, tmp_path): if _slurm_rank_env() is not None: _run_wan22_multinode_slurm_rank(tmp_path) return diff --git a/tests/integration/test_lists/qa/llm_function_multinode.txt b/tests/integration/test_lists/qa/llm_function_multinode.txt index 271f5d295432..cb10040d1934 100644 --- a/tests/integration/test_lists/qa/llm_function_multinode.txt +++ b/tests/integration/test_lists/qa/llm_function_multinode.txt @@ -7,4 +7,3 @@ test_e2e.py::test_multi_nodes_eval[MiniMax-M3-tp16-mmlu] test_e2e.py::test_openai_disagg_multi_nodes_completion[ctx_tp2pp1-gen_tp2pp1] test_e2e.py::test_openai_disagg_multi_nodes_completion[ctx_tp1pp2-gen_tp1pp2] test_e2e.py::test_openai_disagg_multi_nodes_completion_service_discovery[etcd] -examples/visual_gen/test_visual_gen_multi_gpu.py::test_wan22_t2v_lpips_against_golden_multinode_slurm diff --git a/tests/integration/test_lists/test-db/l0_b200_visual_gen_multinode.yml b/tests/integration/test_lists/test-db/l0_b200_visual_gen_multinode.yml index f3cf16d02bb1..667ec7f8ca70 100644 --- a/tests/integration/test_lists/test-db/l0_b200_visual_gen_multinode.yml +++ b/tests/integration/test_lists/test-db/l0_b200_visual_gen_multinode.yml @@ -30,4 +30,4 @@ l0_b200_visual_gen_multinode: stage: post_merge backend: pytorch tests: - - examples/visual_gen/test_visual_gen_multi_gpu.py::test_wan22_t2v_lpips_against_golden_multinode_slurm TIMEOUT (90) + - examples/visual_gen/test_visual_gen_multi_gpu.py::test_wan22_t2v_multinode_slurm_lpips_against_golden TIMEOUT (90) From af11a89a3ca97f62ce9a65467bca17955566ae6a Mon Sep 17 00:00:00 2001 From: Ying Guo <244492186+yingguo-trt@users.noreply.github.com> Date: Mon, 3 Aug 2026 14:54:27 +0800 Subject: [PATCH 12/22] [TRTLLM-13304][test] Register VisualGen multinode case in the QA multinode list The QA cluster pipeline (llm_test_cluster) runs this list with direct srun and skips the trtllm-llmapi-launch wrapper for nodeids containing multinode_slurm, so the SLURM rank environment reaches the test. Register the case so it can be exercised through the QA pipeline before the PR lands. Signed-off-by: Ying Guo <244492186+yingguo-trt@users.noreply.github.com> --- tests/integration/test_lists/qa/llm_function_multinode.txt | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/integration/test_lists/qa/llm_function_multinode.txt b/tests/integration/test_lists/qa/llm_function_multinode.txt index cb10040d1934..05140e1f1ce7 100644 --- a/tests/integration/test_lists/qa/llm_function_multinode.txt +++ b/tests/integration/test_lists/qa/llm_function_multinode.txt @@ -7,3 +7,4 @@ test_e2e.py::test_multi_nodes_eval[MiniMax-M3-tp16-mmlu] test_e2e.py::test_openai_disagg_multi_nodes_completion[ctx_tp2pp1-gen_tp2pp1] test_e2e.py::test_openai_disagg_multi_nodes_completion[ctx_tp1pp2-gen_tp1pp2] test_e2e.py::test_openai_disagg_multi_nodes_completion_service_discovery[etcd] +examples/visual_gen/test_visual_gen_multi_gpu.py::test_wan22_t2v_multinode_slurm_lpips_against_golden From 87a5efe3531b8503ee2d121f3efc88e762056847 Mon Sep 17 00:00:00 2001 From: Ying Guo <244492186+yingguo-trt@users.noreply.github.com> Date: Tue, 4 Aug 2026 09:53:15 +0800 Subject: [PATCH 13/22] [TRTLLM-13304][fix] Capture rank-0 worker thread before shutdown VisualGen.shutdown() drops its executor reference (self.executor = None) after tearing down DiffusionRemoteClient, so reading visual_gen.executor._ext_worker_thread afterwards always resolved to None and the explicit join never ran. pytest-threadleak still flagged the leaked run_diffusion_worker thread. Grab the thread reference before calling shutdown(). The client's own join uses WORKER_TIMEOUT = 2s, which 16-rank dist.destroy_process_group() plus pipeline cleanup routinely outlasts, so the generous follow-up join is what actually lets the thread finish. Signed-off-by: Ying Guo <244492186+yingguo-trt@users.noreply.github.com> --- .../examples/visual_gen/test_visual_gen_multi_gpu.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/tests/integration/defs/examples/visual_gen/test_visual_gen_multi_gpu.py b/tests/integration/defs/examples/visual_gen/test_visual_gen_multi_gpu.py index c1f69adbe6ae..379a2cff17e8 100644 --- a/tests/integration/defs/examples/visual_gen/test_visual_gen_multi_gpu.py +++ b/tests/integration/defs/examples/visual_gen/test_visual_gen_multi_gpu.py @@ -339,11 +339,12 @@ def _run_wan22_multinode_slurm_rank(tmp_path): _assert_lpips_below_threshold(score, WAN_MULTI_GPU_LPIPS_THRESHOLD) finally: if visual_gen is not None: - visual_gen.shutdown() - # shutdown() joins rank 0's in-client worker thread with a bounded - # timeout; 16-rank dist teardown can outlast it and trip - # pytest-threadleak, so wait for the thread explicitly. + # shutdown() joins rank 0's in-client worker thread with a 2s + # timeout and then drops the executor reference, so grab the + # thread first: 16-rank dist teardown outlasts 2s and would + # otherwise trip pytest-threadleak. worker_thread = getattr(visual_gen.executor, "_ext_worker_thread", None) + visual_gen.shutdown() if worker_thread is not None and worker_thread.is_alive(): worker_thread.join(timeout=120) _cleanup_cuda() From ea1ae893f8695828bf8d3e443464ce146019bfc7 Mon Sep 17 00:00:00 2001 From: Ying Guo <244492186+yingguo-trt@users.noreply.github.com> Date: Tue, 4 Aug 2026 11:39:31 +0800 Subject: [PATCH 14/22] [TRTLLM-13304][fix] Install media deps on rank 0 inside the srun step The parent-side _visual_gen_deps call never reached the ranks: each srun step starts a fresh container from the image, so the parent's apt-get install ffmpeg lives only in its own writable layer. rank 0 then hit "MP4 format requires ffmpeg to be installed" when encoding the generated video, because _check_ffmpeg_available() probes the ffmpeg CLI on PATH. Move the fixture call into the rank path, gated on rank == 0 (the only rank that encodes a video), so the deps land in the process that needs them without 16 tasks racing the same apt lock. Signed-off-by: Ying Guo <244492186+yingguo-trt@users.noreply.github.com> --- .../visual_gen/test_visual_gen_multi_gpu.py | 20 ++++++++++++------- 1 file changed, 13 insertions(+), 7 deletions(-) diff --git a/tests/integration/defs/examples/visual_gen/test_visual_gen_multi_gpu.py b/tests/integration/defs/examples/visual_gen/test_visual_gen_multi_gpu.py index 379a2cff17e8..175dec7f5cf9 100644 --- a/tests/integration/defs/examples/visual_gen/test_visual_gen_multi_gpu.py +++ b/tests/integration/defs/examples/visual_gen/test_visual_gen_multi_gpu.py @@ -260,7 +260,7 @@ def _ensure_slurm_external_launch_env(): os.environ.pop(var, None) -def _run_wan22_multinode_slurm_rank(tmp_path): +def _run_wan22_multinode_slurm_rank(request, tmp_path): rank_env = _slurm_rank_env() if rank_env is None: pytest.skip("This VisualGen multi-node case must run under SLURM rank env") @@ -273,6 +273,15 @@ def _run_wan22_multinode_slurm_rank(tmp_path): f"Requires at least {WAN22_LPIPS_MULTINODE_NODES} SLURM nodes, got {node_count}" ) + if rank == 0: + # Only rank 0 encodes the generated video, and only it needs the media + # deps. Installing here (not parent-side) is required: each srun step + # gets a fresh container from the image, so the parent's apt-get never + # reaches the ranks. Keeping it rank-0-only also avoids 16 tasks racing + # the same apt lock; the other ranks simply block in + # dist.init_process_group meanwhile. + request.getfixturevalue("_visual_gen_deps") + _ensure_slurm_external_launch_env() _disable_inductor_compile_worker_quiesce() _parallel_config(**WAN22_LPIPS_MULTINODE_PARALLEL).validate_world_size(world_size) @@ -536,15 +545,12 @@ def test_wan22_t2v_lpips_against_golden_tp(_visual_gen_deps, tmp_path, variant_n def test_wan22_t2v_multinode_slurm_lpips_against_golden(request, tmp_path): if _slurm_rank_env() is not None: - _run_wan22_multinode_slurm_rank(tmp_path) + _run_wan22_multinode_slurm_rank(request, tmp_path) return if os.environ.get(_MULTINODE_SLURM_CHILD_ENV): pytest.skip("VisualGen SLURM child was not launched with SLURM rank env") - # Media deps (av / ffmpeg) are installed parent-side only: rank 0 shares - # node 0's venv with the parent, remote-node workers never touch media, - # and the bare child pytest invocation has no llm_venv harness to run - # the fixture (16 ranks racing pip/apt would be unsafe anyway). - request.getfixturevalue("_visual_gen_deps") + # Media deps are installed by rank 0 inside the srun step, not here: the + # parent only shells out to srun and never encodes a video. _run_wan22_multinode_slurm_parent() From a5bd4eace4df6b515579843739e13425b86b4188 Mon Sep 17 00:00:00 2001 From: Ying Guo <244492186+yingguo-trt@users.noreply.github.com> Date: Wed, 5 Aug 2026 10:05:45 +0800 Subject: [PATCH 15/22] [TRTLLM-13304][test] Drop VisualGen multinode case from the QA multinode list The case is covered by the dedicated L0 post-merge stage (DGX_B200-16_GPUs-2_Nodes-PyTorch-VisualGen-Post-Merge-1 / l0_b200_visual_gen_multinode), which drives srun directly via jenkins/scripts/slurm_run.sh. The QA cluster pipeline wraps every multinode nodeid in trtllm-llmapi-launch, which strips SLURM_* before user code and is incompatible with VisualGen external-launch mode, so registering it there would need a matching trt_jenkins change to be useful. Keep the L0 registration as the single home for this case. Signed-off-by: Ying Guo <244492186+yingguo-trt@users.noreply.github.com> --- tests/integration/test_lists/qa/llm_function_multinode.txt | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/integration/test_lists/qa/llm_function_multinode.txt b/tests/integration/test_lists/qa/llm_function_multinode.txt index 05140e1f1ce7..cb10040d1934 100644 --- a/tests/integration/test_lists/qa/llm_function_multinode.txt +++ b/tests/integration/test_lists/qa/llm_function_multinode.txt @@ -7,4 +7,3 @@ test_e2e.py::test_multi_nodes_eval[MiniMax-M3-tp16-mmlu] test_e2e.py::test_openai_disagg_multi_nodes_completion[ctx_tp2pp1-gen_tp2pp1] test_e2e.py::test_openai_disagg_multi_nodes_completion[ctx_tp1pp2-gen_tp1pp2] test_e2e.py::test_openai_disagg_multi_nodes_completion_service_discovery[etcd] -examples/visual_gen/test_visual_gen_multi_gpu.py::test_wan22_t2v_multinode_slurm_lpips_against_golden From 8cb0f941e74c3e5200dbe49920140d3d8697166a Mon Sep 17 00:00:00 2001 From: Ying Guo <244492186+yingguo-trt@users.noreply.github.com> Date: Wed, 5 Aug 2026 10:09:49 +0800 Subject: [PATCH 16/22] [TRTLLM-13304][chore] Trim test-db YAML header to the two-line SPDX form Signed-off-by: Ying Guo <244492186+yingguo-trt@users.noreply.github.com> --- .../test-db/l0_b200_visual_gen_multinode.yml | 12 ------------ 1 file changed, 12 deletions(-) diff --git a/tests/integration/test_lists/test-db/l0_b200_visual_gen_multinode.yml b/tests/integration/test_lists/test-db/l0_b200_visual_gen_multinode.yml index 667ec7f8ca70..791d909472ae 100644 --- a/tests/integration/test_lists/test-db/l0_b200_visual_gen_multinode.yml +++ b/tests/integration/test_lists/test-db/l0_b200_visual_gen_multinode.yml @@ -1,17 +1,5 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. version: 0.0.1 l0_b200_visual_gen_multinode: From 03617589d32b9afe9db828e73dcf73f891a3aa50 Mon Sep 17 00:00:00 2001 From: Ying Guo <244492186+yingguo-trt@users.noreply.github.com> Date: Wed, 5 Aug 2026 10:33:54 +0800 Subject: [PATCH 17/22] [TRTLLM-13304][test] Rename test-db list to the multi_nodes convention All 22 existing multi-node test-db lists spell it multi_nodes (l0_gb200_multi_nodes.yml, l0_b200_multi_nodes_perf_sanity_*.yml, ...); this file was the only one using "multinode". The difference is not cosmetic: CODEOWNERS line 61 routes tests/integration/test_lists/test-db/*multi_node* to @NVIDIA/trt-llm-multi-gpu-ci-review, and "multinode" does not match that glob, so the new 16-GPU 2-node list silently bypassed that review. Rename the file and its list key, and update the three references in L0_Test.groovy and L0_MergeRequest.groovy. The stage name DGX_B200-16_GPUs-2_Nodes-PyTorch-VisualGen-Post-Merge-1 already follows the existing *-2_Nodes-* convention and is unchanged. Signed-off-by: Ying Guo <244492186+yingguo-trt@users.noreply.github.com> --- jenkins/L0_MergeRequest.groovy | 2 +- jenkins/L0_Test.groovy | 4 ++-- ...l_gen_multinode.yml => l0_b200_visual_gen_multi_nodes.yml} | 2 +- 3 files changed, 4 insertions(+), 4 deletions(-) rename tests/integration/test_lists/test-db/{l0_b200_visual_gen_multinode.yml => l0_b200_visual_gen_multi_nodes.yml} (94%) diff --git a/jenkins/L0_MergeRequest.groovy b/jenkins/L0_MergeRequest.groovy index 155269180596..7f8cae07c105 100644 --- a/jenkins/L0_MergeRequest.groovy +++ b/jenkins/L0_MergeRequest.groovy @@ -1109,7 +1109,7 @@ def getMultiGpuFileChanged(pipeline, testFilter, globalVars) "tests/integration/defs/cpp/test_multi_gpu.py", "tests/integration/test_lists/test-db/l0_b200_multi_gpus_perf_sanity.yml", "tests/integration/test_lists/test-db/l0_b200_multi_nodes_perf_sanity_ctx1_node1_gpu4_gen1_node1_gpu8.yml", - "tests/integration/test_lists/test-db/l0_b200_visual_gen_multinode.yml", + "tests/integration/test_lists/test-db/l0_b200_visual_gen_multi_nodes.yml", "tests/integration/test_lists/test-db/l0_b200_visual_gen_perf_sanity.yml", "tests/integration/test_lists/test-db/l0_dgx_b200.yml", "tests/integration/test_lists/test-db/l0_dgx_b300.yml", diff --git a/jenkins/L0_Test.groovy b/jenkins/L0_Test.groovy index 7cdf229f6ea4..f70644246075 100644 --- a/jenkins/L0_Test.groovy +++ b/jenkins/L0_Test.groovy @@ -1681,7 +1681,7 @@ def runLLMTestlistWithSbatch(pipeline, platform, testList, config=VANILLA_CONFIG // Generate Pytest command String pytestUtil = "" - def visualGenMultinodeSlurmMode = testList == "l0_b200_visual_gen_multinode" + def visualGenMultinodeSlurmMode = testList == "l0_b200_visual_gen_multi_nodes" if (nodeCount > 1 && !visualGenMultinodeSlurmMode) { pytestUtil = "$llmSrcNode/tensorrt_llm/llmapi/trtllm-llmapi-launch" } @@ -5316,7 +5316,7 @@ def launchTestJobs(pipeline, testFilter) "DGX_B300-4_GPUs-PyTorch-Post-Merge-2": ["auto:dgx-b300-flex", "l0_dgx_b300", 2, 2, 4, 1, true], // VisualGen PerfSanity post-merge test "DGX_B200-8_GPUs-PyTorch-VisualGen-PerfSanity-Post-Merge-1": ["auto:dgx-b200-flex", "l0_b200_visual_gen_perf_sanity", 1, 1, 8, 1, true], - "DGX_B200-16_GPUs-2_Nodes-PyTorch-VisualGen-Post-Merge-1": ["auto:dgx-b200-flex", "l0_b200_visual_gen_multinode", 1, 1, 16, 2], + "DGX_B200-16_GPUs-2_Nodes-PyTorch-VisualGen-Post-Merge-1": ["auto:dgx-b200-flex", "l0_b200_visual_gen_multi_nodes", 1, 1, 16, 2], // Single-GPU Gemma4 PerfSanity regression gate and baseline "DGX_B200-PyTorch-PerfSanity-1": ["auto:dgx-b200-flex", "l0_b200_perf_sanity", 1, 1, 1, 1, true], // PerfSanity post-merge tests diff --git a/tests/integration/test_lists/test-db/l0_b200_visual_gen_multinode.yml b/tests/integration/test_lists/test-db/l0_b200_visual_gen_multi_nodes.yml similarity index 94% rename from tests/integration/test_lists/test-db/l0_b200_visual_gen_multinode.yml rename to tests/integration/test_lists/test-db/l0_b200_visual_gen_multi_nodes.yml index 791d909472ae..75cbc59b0e35 100644 --- a/tests/integration/test_lists/test-db/l0_b200_visual_gen_multinode.yml +++ b/tests/integration/test_lists/test-db/l0_b200_visual_gen_multi_nodes.yml @@ -2,7 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 version: 0.0.1 -l0_b200_visual_gen_multinode: +l0_b200_visual_gen_multi_nodes: - condition: ranges: # 2 nodes with each node has 8 GPUs. From 4b42a01d4814e40190063b54903760dc2c575750 Mon Sep 17 00:00:00 2001 From: Ying Guo <244492186+yingguo-trt@users.noreply.github.com> Date: Wed, 5 Aug 2026 10:40:39 +0800 Subject: [PATCH 18/22] [TRTLLM-13304][test] Drop redundant ring rank-linearization test Ring is already covered three ways in this file: the ring+attn2d rejection case, the cfg2/ring2/ulysses2 fiber-flattening case, and TestMultiGPU::test_cfg2_ring2_ulysses2, which validates ring ranks and groups against a real DeviceMesh. Going from ring_size=2 at world_size 8 to ring_size=4 at world_size 16 exercises the same modulo arithmetic and adds no new failure mode. Ring is also not the topology TRTLLM-13304 covers: the multi-node E2E runs cfg x attn2d x ulysses. Keep only the world_size-16 attn2d case, which mirrors that E2E and has no existing counterpart (the current tests top out at 8 GPUs and three dimensions). Signed-off-by: Ying Guo <244492186+yingguo-trt@users.noreply.github.com> --- .../multi_gpu/test_visual_gen_mapping.py | 50 ------------------- 1 file changed, 50 deletions(-) diff --git a/tests/unittest/_torch/visual_gen/multi_gpu/test_visual_gen_mapping.py b/tests/unittest/_torch/visual_gen/multi_gpu/test_visual_gen_mapping.py index 725b0bf7aff1..287d65636271 100644 --- a/tests/unittest/_torch/visual_gen/multi_gpu/test_visual_gen_mapping.py +++ b/tests/unittest/_torch/visual_gen/multi_gpu/test_visual_gen_mapping.py @@ -401,56 +401,6 @@ def test_world_size_16_cfg_attn2d_ulysses_rank_groups(self): DeviceMeshTopologyImpl.device_mesh = old_mesh VisualGenMapping.seq_mesh = old_seq_mesh - def test_world_size_16_cfg_ring_ulysses_rank_groups(self): - from tensorrt_llm._torch.device_mesh import DeviceMeshTopologyImpl - - old_mesh = DeviceMeshTopologyImpl.device_mesh - old_seq_mesh = VisualGenMapping.seq_mesh - try: - for rank in range(16): - DeviceMeshTopologyImpl.device_mesh = None - VisualGenMapping.seq_mesh = None - vgm = VisualGenMapping( - world_size=16, - rank=rank, - cfg_size=2, - ring_size=4, - ulysses_size=2, - ) - - fake_mesh = _FakeDeviceMesh(vgm._dim_names, vgm._dim_sizes, rank) - DeviceMeshTopologyImpl.device_mesh = fake_mesh - VisualGenMapping.seq_mesh = fake_mesh["cp", "ulysses"]._flatten(mesh_dim_name="seq") - - expected_cfg_rank = rank // 8 - expected_cp_rank = (rank // 2) % 4 - expected_ulysses_rank = rank % 2 - expected_seq_rank = expected_cp_rank * 2 + expected_ulysses_rank - - assert vgm.cfg_rank == expected_cfg_rank - assert vgm.tp_rank == 0 - assert vgm.cp_rank == expected_cp_rank - assert vgm.ring_rank == expected_cp_rank - assert vgm.ulysses_rank == expected_ulysses_rank - assert vgm.seq_rank == expected_seq_rank - assert vgm.seq_size == 8 - assert vgm.is_cfg_conditional == (rank < 8) - - cfg_pair_start = rank % 8 - ulysses_pair_start = rank - expected_ulysses_rank - seq_group_start = expected_cfg_rank * 8 - cp_group_start = expected_cfg_rank * 8 + expected_ulysses_rank - - assert vgm.cfg_group == (cfg_pair_start, cfg_pair_start + 8) - assert vgm.tp_group_pg == (rank,) - assert vgm.ulysses_group == (ulysses_pair_start, ulysses_pair_start + 1) - assert vgm.cp_group == tuple(cp_group_start + stride for stride in (0, 2, 4, 6)) - assert vgm.ring_group == vgm.cp_group - assert vgm.seq_group() == tuple(range(seq_group_start, seq_group_start + 8)) - finally: - DeviceMeshTopologyImpl.device_mesh = old_mesh - VisualGenMapping.seq_mesh = old_seq_mesh - # ============================================================================= # Multi-GPU tests — validate actual DeviceMesh groups and ranks From 473bb40a7acbb672fe73364422e6a5c92fa75734 Mon Sep 17 00:00:00 2001 From: Ying Guo <244492186+yingguo-trt@users.noreply.github.com> Date: Wed, 5 Aug 2026 10:55:39 +0800 Subject: [PATCH 19/22] [TRTLLM-13304][test] Address review feedback on the multinode case - slurm_run.sh: bound the rank-0 install-marker wait with a deadline (TRTLLM_VISUAL_GEN_MULTINODE_INSTALL_TIMEOUT, default 3600s) and print the markers found, instead of looping forever if a node never reports. - Parent path: drive srun through defs.trt_test_alternative.popen so the whole process tree is killed on timeout; plain subprocess.run only reaps the local srun and can strand remote tasks. popen also refills TimeoutExpired.output, so the failure message keeps the srun log. - Parent path: validate the per-rank pytest summaries that srun -l emits (exactly one "1 passed" per rank, ANSI stripped) rather than searching the whole log for "passed"/"skipped", which any log line could satisfy. - Cross-reference WAN22_LPIPS_MULTINODE_PARALLEL and the 16-rank mapping unit test in both directions; nothing enforces that they agree, and a drifted unit test would silently guard a topology nobody runs. Signed-off-by: Ying Guo <244492186+yingguo-trt@users.noreply.github.com> --- jenkins/scripts/slurm_run.sh | 16 ++++- .../visual_gen/test_visual_gen_multi_gpu.py | 59 +++++++++---------- .../multi_gpu/test_visual_gen_mapping.py | 10 ++++ 3 files changed, 51 insertions(+), 34 deletions(-) diff --git a/jenkins/scripts/slurm_run.sh b/jenkins/scripts/slurm_run.sh index e5791375349f..64c4c124992d 100755 --- a/jenkins/scripts/slurm_run.sh +++ b/jenkins/scripts/slurm_run.sh @@ -65,8 +65,20 @@ if [[ "${TRTLLM_VISUAL_GEN_MULTINODE_SLURM_PARENT:-0}" == "1" ]]; then if [[ "${SLURM_PROCID:-0}" == "0" ]]; then expected_nodes="${SLURM_JOB_NUM_NODES:-${SLURM_NNODES:-1}}" - while [[ "$(find "$install_done_dir" -maxdepth 1 -type f -name 'node_*' | wc -l)" -lt "$expected_nodes" ]]; do - echo "Waiting for VisualGen multi-node install markers in $install_done_dir" + install_wait_timeout="${TRTLLM_VISUAL_GEN_MULTINODE_INSTALL_TIMEOUT:-3600}" + install_wait_deadline=$((SECONDS + install_wait_timeout)) + while true; do + marker_count="$(find "$install_done_dir" -maxdepth 1 -type f -name 'node_*' -print | wc -l)" + if [[ "$marker_count" -ge "$expected_nodes" ]]; then + break + fi + if ((SECONDS >= install_wait_deadline)); then + echo "Timed out after ${install_wait_timeout}s waiting for VisualGen install markers (${marker_count}/${expected_nodes})." + echo "Markers present in $install_done_dir:" + find "$install_done_dir" -maxdepth 1 -type f -name 'node_*' -print | sort || true + exit 1 + fi + echo "Waiting for VisualGen multi-node install markers in $install_done_dir (${marker_count}/${expected_nodes})" sleep 10 done pytestCommand="env -u SLURM_PROCID -u SLURM_NTASKS -u SLURM_LOCALID -u SLURM_NODEID -u SLURM_GTIDS ${pytestCommand}" diff --git a/tests/integration/defs/examples/visual_gen/test_visual_gen_multi_gpu.py b/tests/integration/defs/examples/visual_gen/test_visual_gen_multi_gpu.py index 175dec7f5cf9..8e075fe517aa 100644 --- a/tests/integration/defs/examples/visual_gen/test_visual_gen_multi_gpu.py +++ b/tests/integration/defs/examples/visual_gen/test_visual_gen_multi_gpu.py @@ -16,6 +16,7 @@ import glob import os +import re import shutil import subprocess import sys @@ -47,6 +48,7 @@ _save_lpips_video_mp4, _skip_if_missing, ) +from defs.trt_test_alternative import popen def _parallel_config(**kwargs): @@ -82,6 +84,7 @@ def _parallel_config(**kwargs): WAN22_LPIPS_MULTINODE_WORLD_SIZE = 16 WAN22_LPIPS_MULTINODE_NODES = 2 WAN22_LPIPS_MULTINODE_GPUS_PER_NODE = 8 +# Keep this topology aligned with its 16-rank mapping unit test. WAN22_LPIPS_MULTINODE_PARALLEL = { "cfg_size": 2, "attn2d_size": (2, 2), @@ -218,8 +221,7 @@ def _trtllm_launch_wrapper_world_size(): def _multinode_subprocess_timeout(): - # Default below the QA pipeline's outer --timeout=3600s so the inner srun - # timeout (with captured output) fires before pytest-timeout kills the test. + # Leave headroom for the outer 3600-second pytest timeout. return int(os.environ.get("TRTLLM_VISUAL_GEN_MULTINODE_TIMEOUT", "3300")) @@ -254,8 +256,7 @@ def _ensure_slurm_external_launch_env(): os.environ["MASTER_ADDR"] = _resolve_slurm_master_addr() os.environ.setdefault("MASTER_PORT", _default_master_port()) - # Force VisualGen to exercise its SLURM detection branch even when the - # surrounding CI wrapper leaves torchrun-like variables behind. + # Prefer SLURM ranks over inherited torchrun variables. for var in ("RANK", "WORLD_SIZE", "LOCAL_RANK"): os.environ.pop(var, None) @@ -274,12 +275,7 @@ def _run_wan22_multinode_slurm_rank(request, tmp_path): ) if rank == 0: - # Only rank 0 encodes the generated video, and only it needs the media - # deps. Installing here (not parent-side) is required: each srun step - # gets a fresh container from the image, so the parent's apt-get never - # reaches the ranks. Keeping it rank-0-only also avoids 16 tasks racing - # the same apt lock; the other ranks simply block in - # dist.init_process_group meanwhile. + # Each srun rank gets a fresh container; rank 0 installs deps to avoid apt-lock races. request.getfixturevalue("_visual_gen_deps") _ensure_slurm_external_launch_env() @@ -348,10 +344,7 @@ def _run_wan22_multinode_slurm_rank(request, tmp_path): _assert_lpips_below_threshold(score, WAN_MULTI_GPU_LPIPS_THRESHOLD) finally: if visual_gen is not None: - # shutdown() joins rank 0's in-client worker thread with a 2s - # timeout and then drops the executor reference, so grab the - # thread first: 16-rank dist teardown outlasts 2s and would - # otherwise trip pytest-threadleak. + # Capture before shutdown drops the executor; its 2s join is too short here. worker_thread = getattr(visual_gen.executor, "_ext_worker_thread", None) visual_gen.shutdown() if worker_thread is not None and worker_thread.is_alive(): @@ -381,9 +374,7 @@ def _run_wan22_multinode_slurm_parent(): if shutil.which("srun") is None: pytest.skip("srun is required for the VisualGen multi-node LPIPS case") - # Pre-check the checkpoint here so a missing model skips the parent honestly, - # instead of letting every srun rank skip and the parent report a false pass - # (an all-skipped pytest run still exits 0). + # Avoid treating an all-skipped child run as a pass. _skip_if_missing( _lpips_model_path("Wan2.2-T2V-A14B-Diffusers"), "Wan 2.2 checkpoint", @@ -416,16 +407,16 @@ def _run_wan22_multinode_slurm_parent(): nodeid, ] try: - result = subprocess.run( + with popen( cmd, cwd=Path(__file__).resolve().parents[5], env=env, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, - check=False, - timeout=_multinode_subprocess_timeout(), - ) + ) as proc: + output, _ = proc.communicate(timeout=_multinode_subprocess_timeout()) + returncode = proc.returncode except subprocess.TimeoutExpired as exc: output = exc.output or "" pytest.fail( @@ -433,20 +424,26 @@ def _run_wan22_multinode_slurm_parent(): f"{exc.timeout} seconds:\n{output}" ) - if result.returncode != 0: + if returncode != 0: pytest.fail( - "VisualGen multi-node SLURM subprocess failed with " - f"exit code {result.returncode}:\n{result.stdout}" + f"VisualGen multi-node SLURM subprocess failed with exit code {returncode}:\n{output}" ) - # Guard against a false pass: an all-skipped/empty pytest run also exits 0. - # The parent has already validated every precondition, so once srun launches - # all ranks must actually pass; any skip or empty collection is a real defect. - output = result.stdout or "" - if "passed" not in output or "skipped" in output or "no tests ran" in output: + # Validate rank-prefixed pytest summaries, not incidental log text. + summary_output = re.sub(r"\x1b\[[0-9;]*m", "", output) + rank_summaries = sorted( + (int(rank), int(count), outcome) + for rank, count, outcome in re.findall( + r"(?m)^\s*(\d+):\s+(?:=+\s*)?(\d+)\s+(passed|skipped)\b", + summary_output, + ) + ) + expected_summaries = [(rank, 1, "passed") for rank in range(WAN22_LPIPS_MULTINODE_WORLD_SIZE)] + if rank_summaries != expected_summaries: pytest.fail( "VisualGen multi-node SLURM run exited 0 but did not actually execute " - "(expected all ranks to pass with no skips):\n" + "(expected exactly one passing pytest summary for every rank):\n" + f"observed summaries: {rank_summaries}\n" f"{output}" ) @@ -551,6 +548,4 @@ def test_wan22_t2v_multinode_slurm_lpips_against_golden(request, tmp_path): if os.environ.get(_MULTINODE_SLURM_CHILD_ENV): pytest.skip("VisualGen SLURM child was not launched with SLURM rank env") - # Media deps are installed by rank 0 inside the srun step, not here: the - # parent only shells out to srun and never encodes a video. _run_wan22_multinode_slurm_parent() diff --git a/tests/unittest/_torch/visual_gen/multi_gpu/test_visual_gen_mapping.py b/tests/unittest/_torch/visual_gen/multi_gpu/test_visual_gen_mapping.py index 287d65636271..6b573bff1f88 100644 --- a/tests/unittest/_torch/visual_gen/multi_gpu/test_visual_gen_mapping.py +++ b/tests/unittest/_torch/visual_gen/multi_gpu/test_visual_gen_mapping.py @@ -337,6 +337,16 @@ def test_mixed_parallelism(self): class TestRankLinearizationWithoutDist: def test_world_size_16_cfg_attn2d_ulysses_rank_groups(self): + """CPU-side guard for the 16-GPU multi-node E2E's rank linearization. + + The topology below must stay in sync with + WAN22_LPIPS_MULTINODE_PARALLEL in + tests/integration/defs/examples/visual_gen/test_visual_gen_multi_gpu.py. + Nothing enforces that: the two trees have separate conftests, so + importing the constant would drag integration fixtures into a + CPU-only unit test. If you change the E2E's parallel config, change + this test too, or it silently guards a topology nobody runs. + """ from tensorrt_llm._torch.device_mesh import DeviceMeshTopologyImpl old_mesh = DeviceMeshTopologyImpl.device_mesh From 18b6dff27c12038e48baeaf939885072f2a10c70 Mon Sep 17 00:00:00 2001 From: Ying Guo <244492186+yingguo-trt@users.noreply.github.com> Date: Wed, 5 Aug 2026 11:31:48 +0800 Subject: [PATCH 20/22] [TRTLLM-13304][test] Narrow the PR to the multinode E2E and fix two guards Scope: drop both unit-test additions. This PR does not modify VisualGenMapping or the external-launch code they guard, real multi-GPU mapping tests already cover attn2d/ulysses grouping against a live DeviceMesh, and the fake DeviceMesh re-implemented init_device_mesh's row-major rule rather than verifying it. Both files are now byte-identical to main; the PR is the multi-node E2E plus its CI registration. Fixes: - Guard on TLLM_SPAWN_PROXY_PROCESS alone. trtllm-llmapi-launch strips SLURM_* at any world size, so gating on tllm_mpi_size >= 16 let the stripped-env case fall through to a "no SLURM allocation" skip that points at the wrong cause. Drops the now-unused helper. - Scope the install markers to ${SLURM_JOB_ID}. runLLMTestlistOnSlurm retries on infra failure and can reuse jobWorkspace; stale node_* files would let rank 0 stop waiting before this attempt's nodes finish installing, resurfacing the missing-ffmpeg failure intermittently. Signed-off-by: Ying Guo <244492186+yingguo-trt@users.noreply.github.com> --- jenkins/scripts/slurm_run.sh | 5 +- .../visual_gen/test_visual_gen_multi_gpu.py | 16 +- .../multi_gpu/test_visual_gen_mapping.py | 139 ------------------ .../multi_gpu/test_visual_gen_multinode.py | 100 +------------ 4 files changed, 9 insertions(+), 251 deletions(-) diff --git a/jenkins/scripts/slurm_run.sh b/jenkins/scripts/slurm_run.sh index 64c4c124992d..cbcbbca27ac6 100755 --- a/jenkins/scripts/slurm_run.sh +++ b/jenkins/scripts/slurm_run.sh @@ -57,7 +57,10 @@ env | sort echo "Full Command: $pytestCommand" if [[ "${TRTLLM_VISUAL_GEN_MULTINODE_SLURM_PARENT:-0}" == "1" ]]; then - install_done_dir="${jobWorkspace}/visual_gen_multinode_install_done" + # Scope the markers to the SLURM job: the stage's infra-retry loop can reuse + # jobWorkspace, and stale node_* files from a previous attempt would let + # rank 0 start before this attempt's nodes finish installing. + install_done_dir="${jobWorkspace}/visual_gen_multinode_install_done_${SLURM_JOB_ID:-0}" mkdir -p "$install_done_dir" if [[ "${SLURM_LOCALID:-0}" == "0" ]]; then touch "${install_done_dir}/node_${SLURM_NODEID:-0}" diff --git a/tests/integration/defs/examples/visual_gen/test_visual_gen_multi_gpu.py b/tests/integration/defs/examples/visual_gen/test_visual_gen_multi_gpu.py index 8e075fe517aa..9ad6338191c3 100644 --- a/tests/integration/defs/examples/visual_gen/test_visual_gen_multi_gpu.py +++ b/tests/integration/defs/examples/visual_gen/test_visual_gen_multi_gpu.py @@ -84,7 +84,6 @@ def _parallel_config(**kwargs): WAN22_LPIPS_MULTINODE_WORLD_SIZE = 16 WAN22_LPIPS_MULTINODE_NODES = 2 WAN22_LPIPS_MULTINODE_GPUS_PER_NODE = 8 -# Keep this topology aligned with its 16-rank mapping unit test. WAN22_LPIPS_MULTINODE_PARALLEL = { "cfg_size": 2, "attn2d_size": (2, 2), @@ -213,13 +212,6 @@ def _slurm_node_count(): return None -def _trtllm_launch_wrapper_world_size(): - try: - return int(os.environ.get("tllm_mpi_size", "1") or 1) - except ValueError: - return 1 - - def _multinode_subprocess_timeout(): # Leave headroom for the outer 3600-second pytest timeout. return int(os.environ.get("TRTLLM_VISUAL_GEN_MULTINODE_TIMEOUT", "3300")) @@ -353,10 +345,10 @@ def _run_wan22_multinode_slurm_rank(request, tmp_path): def _run_wan22_multinode_slurm_parent(): - if ( - os.environ.get("TLLM_SPAWN_PROXY_PROCESS") == "1" - and _trtllm_launch_wrapper_world_size() >= WAN22_LPIPS_MULTINODE_WORLD_SIZE - ): + # Guard on the wrapper alone, not on its world size: trtllm-llmapi-launch + # strips SLURM_* at any scale, and a size-gated check would let the + # stripped-env case fall through to a misleading "no SLURM allocation" skip. + if os.environ.get("TLLM_SPAWN_PROXY_PROCESS") == "1": pytest.fail( "VisualGen SLURM external-launch coverage cannot run under " "trtllm-llmapi-launch because that wrapper removes SLURM_* env " diff --git a/tests/unittest/_torch/visual_gen/multi_gpu/test_visual_gen_mapping.py b/tests/unittest/_torch/visual_gen/multi_gpu/test_visual_gen_mapping.py index 6b573bff1f88..c73762ae3c7a 100644 --- a/tests/unittest/_torch/visual_gen/multi_gpu/test_visual_gen_mapping.py +++ b/tests/unittest/_torch/visual_gen/multi_gpu/test_visual_gen_mapping.py @@ -72,68 +72,6 @@ def _run_multi_gpu(world_size, test_fn): ) -def _row_major_coords(rank, dim_names, dim_sizes): - # Mirrors torch.distributed.device_mesh.init_device_mesh row-major layout; - # the existing multi-GPU mapping tests anchor this assumption against real DeviceMesh. - coords = {} - remaining = rank - for dim in reversed(dim_names): - size = dim_sizes[dim] - coords[dim] = remaining % size - remaining //= size - return coords - - -class _FakeDeviceMeshView: - def __init__(self, dim_names, dim_sizes, rank, group_dims): - self._dim_names = dim_names - self._dim_sizes = dim_sizes - self._rank = rank - self._group_dims = tuple(group_dims) - self._coords = _row_major_coords(rank, dim_names, dim_sizes) - self._world_size = 1 - for size in dim_sizes.values(): - self._world_size *= size - - def __getitem__(self, dim): - if isinstance(dim, tuple): - group_dims = dim - else: - group_dims = (dim,) - return _FakeDeviceMeshView( - self._dim_names, - self._dim_sizes, - self._rank, - group_dims, - ) - - def _flatten(self, mesh_dim_name): - return _FakeDeviceMeshView( - self._dim_names, - self._dim_sizes, - self._rank, - self._group_dims, - ) - - def get_local_rank(self): - assert len(self._group_dims) == 1 - return self._coords[self._group_dims[0]] - - def get_group(self): - fixed_dims = set(self._dim_names) - set(self._group_dims) - group = [] - for rank in range(self._world_size): - coords = _row_major_coords(rank, self._dim_names, self._dim_sizes) - if all(coords[dim] == self._coords[dim] for dim in fixed_dims): - group.append(rank) - return tuple(group) - - -class _FakeDeviceMesh(_FakeDeviceMeshView): - def __init__(self, dim_names, dim_sizes, rank): - super().__init__(dim_names, dim_sizes, rank, dim_names) - - # ============================================================================= # Single-GPU tests (no dist required) # ============================================================================= @@ -335,83 +273,6 @@ def test_mixed_parallelism(self): assert m.world_size == 2 -class TestRankLinearizationWithoutDist: - def test_world_size_16_cfg_attn2d_ulysses_rank_groups(self): - """CPU-side guard for the 16-GPU multi-node E2E's rank linearization. - - The topology below must stay in sync with - WAN22_LPIPS_MULTINODE_PARALLEL in - tests/integration/defs/examples/visual_gen/test_visual_gen_multi_gpu.py. - Nothing enforces that: the two trees have separate conftests, so - importing the constant would drag integration fixtures into a - CPU-only unit test. If you change the E2E's parallel config, change - this test too, or it silently guards a topology nobody runs. - """ - from tensorrt_llm._torch.device_mesh import DeviceMeshTopologyImpl - - old_mesh = DeviceMeshTopologyImpl.device_mesh - old_seq_mesh = VisualGenMapping.seq_mesh - try: - for rank in range(16): - DeviceMeshTopologyImpl.device_mesh = None - VisualGenMapping.seq_mesh = None - vgm = VisualGenMapping( - world_size=16, - rank=rank, - cfg_size=2, - attn2d_row_size=2, - attn2d_col_size=2, - ulysses_size=2, - ) - - fake_mesh = _FakeDeviceMesh(vgm._dim_names, vgm._dim_sizes, rank) - DeviceMeshTopologyImpl.device_mesh = fake_mesh - VisualGenMapping.seq_mesh = fake_mesh["cp_row", "cp_col", "ulysses"]._flatten( - mesh_dim_name="seq" - ) - vgm._attach_attn2d_groups_from_device_mesh() - - expected_cfg_rank = rank // 8 - expected_cp_row_rank = (rank // 4) % 2 - expected_cp_col_rank = (rank // 2) % 2 - expected_ulysses_rank = rank % 2 - expected_cp_rank = expected_cp_row_rank * 2 + expected_cp_col_rank - expected_seq_rank = expected_cp_rank * 2 + expected_ulysses_rank - - assert vgm.cfg_rank == expected_cfg_rank - assert vgm.tp_rank == 0 - assert vgm.cp_row_rank == expected_cp_row_rank - assert vgm.cp_col_rank == expected_cp_col_rank - assert vgm.cp_rank == expected_cp_rank - assert vgm.attn2d_mesh_rank == expected_cp_rank - assert vgm.ulysses_rank == expected_ulysses_rank - assert vgm.seq_rank == expected_seq_rank - assert vgm.seq_size == 8 - assert vgm.is_cfg_conditional == (rank < 8) - - cfg_pair_start = rank % 8 - ulysses_pair_start = rank - expected_ulysses_rank - seq_group_start = expected_cfg_rank * 8 - cp_group_start = expected_cfg_rank * 8 + expected_ulysses_rank - row_group_start = ( - expected_cfg_rank * 8 + expected_cp_row_rank * 4 + expected_ulysses_rank - ) - col_group_start = ( - expected_cfg_rank * 8 + expected_cp_col_rank * 2 + expected_ulysses_rank - ) - - assert vgm.cfg_group == (cfg_pair_start, cfg_pair_start + 8) - assert vgm.ulysses_group == (ulysses_pair_start, ulysses_pair_start + 1) - assert vgm.cp_group == tuple(cp_group_start + stride for stride in (0, 2, 4, 6)) - assert vgm.attn2d_mesh_group == vgm.cp_group - assert vgm.attn2d_row_group == (row_group_start, row_group_start + 2) - assert vgm.attn2d_col_group == (col_group_start, col_group_start + 4) - assert vgm.seq_group() == tuple(range(seq_group_start, seq_group_start + 8)) - finally: - DeviceMeshTopologyImpl.device_mesh = old_mesh - VisualGenMapping.seq_mesh = old_seq_mesh - - # ============================================================================= # Multi-GPU tests — validate actual DeviceMesh groups and ranks # ============================================================================= diff --git a/tests/unittest/_torch/visual_gen/multi_gpu/test_visual_gen_multinode.py b/tests/unittest/_torch/visual_gen/multi_gpu/test_visual_gen_multinode.py index 3bdfdec1cffc..fddadb291d60 100644 --- a/tests/unittest/_torch/visual_gen/multi_gpu/test_visual_gen_multinode.py +++ b/tests/unittest/_torch/visual_gen/multi_gpu/test_visual_gen_multinode.py @@ -13,11 +13,7 @@ from tensorrt_llm._torch.visual_gen.executor import run_diffusion_worker from tensorrt_llm.visual_gen.args import VisualGenArgs -from tensorrt_llm.visual_gen.visual_gen import ( - DiffusionRemoteClient, - VisualGen, - _detect_external_launch, -) +from tensorrt_llm.visual_gen.visual_gen import DiffusionRemoteClient, _detect_external_launch # ============================================================================= # _detect_external_launch() @@ -252,97 +248,3 @@ def pre_set_event(): f"Worker {i}: expected local_rank={i}, got {kwargs['local_rank']}. " f"With LOCAL_RANK=0 in env, all workers would get device_id=0." ) - - -class TestVisualGenExternalLaunchInit: - def test_world_size_mismatch_raises(self): - args = VisualGenArgs(model="/tmp/model", parallel_config={"ulysses_size": 4}) - - with patch( - "tensorrt_llm.visual_gen.visual_gen._detect_external_launch", - return_value=(0, 0, 8, "node0", 29500), - ): - with pytest.raises(ValueError, match="does not match"): - VisualGen(model="/tmp/model", args=args) - - def test_nonzero_rank_runs_worker_then_exits(self): - args = VisualGenArgs(model="/tmp/model", parallel_config={"ulysses_size": 4}) - mock_worker = MagicMock() - - with ( - patch( - "tensorrt_llm.visual_gen.visual_gen._detect_external_launch", - return_value=(2, 2, 4, "node0", 29500), - ), - patch("tensorrt_llm.visual_gen.visual_gen.run_diffusion_worker", mock_worker), - ): - with pytest.raises(SystemExit) as exc_info: - VisualGen(model="/tmp/model", args=args) - - assert exc_info.value.code == 0 - kwargs = mock_worker.call_args.kwargs - assert ( - kwargs["rank"], - kwargs["local_rank"], - kwargs["world_size"], - kwargs["master_addr"], - kwargs["master_port"], - ) == (2, 2, 4, "node0", 29500) - assert kwargs["request_queue_addr"] is kwargs["response_queue_addr"] is None - - -class TestExternalLaunchClientWiring: - @pytest.fixture(autouse=True) - def _clean_env(self, monkeypatch): - for var in ["RANK", "WORLD_SIZE", "LOCAL_RANK", "SLURM_PROCID", "SLURM_NTASKS"]: - monkeypatch.delenv(var, raising=False) - - def test_rank0_uses_master_addr_and_thread_worker(self): - args = VisualGenArgs(model="/tmp/model", parallel_config={"ulysses_size": 4}) - - mock_ctx = MagicMock() - original_event = threading.Event - - def pre_set_event(): - event = original_event() - event.set() - return event - - with ( - patch( - "tensorrt_llm._torch.visual_gen.executor._detect_external_launch", - return_value=(0, 0, 4, "node0", 29500), - ), - patch("tensorrt_llm._torch.visual_gen.executor.mp.get_context", return_value=mock_ctx), - patch("tensorrt_llm._torch.visual_gen.executor.threading.Thread") as mock_thread_cls, - patch( - "tensorrt_llm._torch.visual_gen.executor.threading.Event", side_effect=pre_set_event - ), - patch.object(DiffusionRemoteClient, "_wait_ready"), - ): - mock_thread_cls.return_value = MagicMock() - client = DiffusionRemoteClient(args=args) - - assert (client.master_addr, client.master_port) == ("node0", 29500) - assert client.req_addr_connect.startswith("tcp://node0:") - assert client.resp_addr_connect.startswith("tcp://node0:") - mock_ctx.Process.assert_not_called() - worker_calls = [ - c - for c in mock_thread_cls.call_args_list - if "in_client_process" in c.kwargs.get("kwargs", {}) - ] - assert len(worker_calls) == 1 - thread_call = worker_calls[0] - assert thread_call.kwargs["daemon"] is True - worker_kwargs = thread_call.kwargs["kwargs"] - assert worker_kwargs["in_client_process"] - assert ( - worker_kwargs["rank"], - worker_kwargs["local_rank"], - worker_kwargs["world_size"], - worker_kwargs["master_addr"], - worker_kwargs["master_port"], - ) == (0, 0, 4, "node0", 29500) - assert worker_kwargs["request_queue_addr"] == client.req_addr_connect - assert worker_kwargs["response_queue_addr"] == client.resp_addr_connect From 4cb0d10f4ba740cfb965f3b87221905826fa2482 Mon Sep 17 00:00:00 2001 From: Ying Guo <244492186+yingguo-trt@users.noreply.github.com> Date: Wed, 5 Aug 2026 11:38:05 +0800 Subject: [PATCH 21/22] [TRTLLM-13304][ci] Disable PMIx bootstrap for the VisualGen multinode step The outer srun step for this stage runs only a shell script that starts pytest; the real 16-way fan-out is an inner srun the parent issues itself, and VisualGen never initializes MPI (TLLM_DISABLE_MPI=1 + torch.distributed). Bootstrapping PMIx there injects PMIX_* describing the 2-task outer step, and slurm_run.sh only unsets five SLURM_* names, so the inner srun --export=ALL hands those stale variables to all 16 child ranks. Pass --mpi=none rather than dropping the flag: with no --mpi, srun falls back to the cluster's MpiDefault, which is commonly pmix, so omitting it would not reliably disable the bootstrap. Other stages keep --mpi=pmix unchanged. Also trims the two verbose comments added in the previous commit to match the surrounding one-line style. Signed-off-by: Ying Guo <244492186+yingguo-trt@users.noreply.github.com> --- jenkins/L0_Test.groovy | 2 +- jenkins/scripts/slurm_run.sh | 4 +--- .../defs/examples/visual_gen/test_visual_gen_multi_gpu.py | 4 +--- 3 files changed, 3 insertions(+), 7 deletions(-) diff --git a/jenkins/L0_Test.groovy b/jenkins/L0_Test.groovy index f70644246075..3d46ddc2d666 100644 --- a/jenkins/L0_Test.groovy +++ b/jenkins/L0_Test.groovy @@ -1937,7 +1937,7 @@ def runLLMTestlistWithSbatch(pipeline, platform, testList, config=VANILLA_CONFIG """ } else { if(nodeCount > 1) { - srunArgs.add("--mpi=pmix") + srunArgs.add(visualGenMultinodeSlurmMode ? "--mpi=none" : "--mpi=pmix") } def scriptContent = """ diff --git a/jenkins/scripts/slurm_run.sh b/jenkins/scripts/slurm_run.sh index cbcbbca27ac6..cbaea8cddb64 100755 --- a/jenkins/scripts/slurm_run.sh +++ b/jenkins/scripts/slurm_run.sh @@ -57,9 +57,7 @@ env | sort echo "Full Command: $pytestCommand" if [[ "${TRTLLM_VISUAL_GEN_MULTINODE_SLURM_PARENT:-0}" == "1" ]]; then - # Scope the markers to the SLURM job: the stage's infra-retry loop can reuse - # jobWorkspace, and stale node_* files from a previous attempt would let - # rank 0 start before this attempt's nodes finish installing. + # Scope to the job: infra retries reuse jobWorkspace, and stale markers would end the wait early. install_done_dir="${jobWorkspace}/visual_gen_multinode_install_done_${SLURM_JOB_ID:-0}" mkdir -p "$install_done_dir" if [[ "${SLURM_LOCALID:-0}" == "0" ]]; then diff --git a/tests/integration/defs/examples/visual_gen/test_visual_gen_multi_gpu.py b/tests/integration/defs/examples/visual_gen/test_visual_gen_multi_gpu.py index 9ad6338191c3..fb08d4ba712a 100644 --- a/tests/integration/defs/examples/visual_gen/test_visual_gen_multi_gpu.py +++ b/tests/integration/defs/examples/visual_gen/test_visual_gen_multi_gpu.py @@ -345,9 +345,7 @@ def _run_wan22_multinode_slurm_rank(request, tmp_path): def _run_wan22_multinode_slurm_parent(): - # Guard on the wrapper alone, not on its world size: trtllm-llmapi-launch - # strips SLURM_* at any scale, and a size-gated check would let the - # stripped-env case fall through to a misleading "no SLURM allocation" skip. + # The wrapper strips SLURM_* at any scale, so do not gate this on world size. if os.environ.get("TLLM_SPAWN_PROXY_PROCESS") == "1": pytest.fail( "VisualGen SLURM external-launch coverage cannot run under " From 53ce5b2acfcad7064813fc23a1e22f96471a5c97 Mon Sep 17 00:00:00 2001 From: Ying Guo <244492186+yingguo-trt@users.noreply.github.com> Date: Wed, 5 Aug 2026 11:42:01 +0800 Subject: [PATCH 22/22] [TRTLLM-13304][ci] Revert the SLURM_JOB_ID marker suffix The suffix guarded a problem that does not exist. runLLMTestlistWithSbatch derives jobWorkspace from a fresh UUID on every call (L0_Test.groovy:1490) and the infra-retry loop re-enters that function per attempt, so retries never share a marker directory. The submit script also rm -rf's everything outside the keep list immediately before sbatch (L0_Test.groovy:1968). The one path that does reuse a workspace -- "reuse an already-active job" -- reuses the same SLURM job, so SLURM_JOB_ID would be identical and the suffix would protect nothing. Restores the plain marker directory and drops the comment, which claimed infra retries reuse jobWorkspace and was simply wrong. The wrapper-guard fix and --mpi=none stay. Signed-off-by: Ying Guo <244492186+yingguo-trt@users.noreply.github.com> --- jenkins/scripts/slurm_run.sh | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/jenkins/scripts/slurm_run.sh b/jenkins/scripts/slurm_run.sh index cbaea8cddb64..64c4c124992d 100755 --- a/jenkins/scripts/slurm_run.sh +++ b/jenkins/scripts/slurm_run.sh @@ -57,8 +57,7 @@ env | sort echo "Full Command: $pytestCommand" if [[ "${TRTLLM_VISUAL_GEN_MULTINODE_SLURM_PARENT:-0}" == "1" ]]; then - # Scope to the job: infra retries reuse jobWorkspace, and stale markers would end the wait early. - install_done_dir="${jobWorkspace}/visual_gen_multinode_install_done_${SLURM_JOB_ID:-0}" + install_done_dir="${jobWorkspace}/visual_gen_multinode_install_done" mkdir -p "$install_done_dir" if [[ "${SLURM_LOCALID:-0}" == "0" ]]; then touch "${install_done_dir}/node_${SLURM_NODEID:-0}"