From b656fcf45bdbe75df6f47911275dd248c6c5736f Mon Sep 17 00:00:00 2001 From: Nick Cao Date: Mon, 28 Sep 2026 14:22:25 -0400 Subject: [PATCH] [Tests] Use get_file_store_init_method helper from vllm (#7528) Signed-off-by: Nick Cao Co-authored-by: Codex --- tests/diffusion/attention/test_ulysses_uaa.py | 16 ++++++++-------- tests/diffusion/cache/test_seacache.py | 6 +++--- tests/diffusion/distributed/test_hsdp.py | 4 ++-- .../test_parallel_state_sp_groups.py | 4 ++-- .../distributed/test_pipeline_parallel.py | 18 +++++++++--------- .../distributed/test_ulysses_uaa_perf.py | 4 ++-- tests/diffusion/offloader/conftest.py | 4 ++-- tests/helpers/runtime.py | 14 -------------- 8 files changed, 28 insertions(+), 42 deletions(-) diff --git a/tests/diffusion/attention/test_ulysses_uaa.py b/tests/diffusion/attention/test_ulysses_uaa.py index 3782d8dd6..f97d26eec 100644 --- a/tests/diffusion/attention/test_ulysses_uaa.py +++ b/tests/diffusion/attention/test_ulysses_uaa.py @@ -11,9 +11,9 @@ from types import SimpleNamespace import numpy as np import pytest import torch +from vllm.utils.network_utils import get_file_store_init_method from tests.helpers.mark import hardware_test -from tests.helpers.runtime import get_distributed_init_method from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata from vllm_omni.diffusion.attention.layer import Attention from vllm_omni.diffusion.attention.parallel.ulysses import UlyssesParallelAttention @@ -783,8 +783,8 @@ def test_ulysses_uaa_matches_baseline( kv_heads = num_kv_heads if num_kv_heads is not None else num_heads has_joint = joint_len is not None - base_init_method = get_distributed_init_method() - sp_init_method = get_distributed_init_method() + base_init_method = get_file_store_init_method() + sp_init_method = get_file_store_init_method() with tempfile.NamedTemporaryFile(delete=False, suffix=".npz") as f_in: input_file = f_in.name @@ -989,7 +989,7 @@ def test_ulysses_uaa_contract_skips_length_gather_matches_slow_path(num_heads: i head_size = 8 seq_len = 6 # divisible by 2 so both ranks hold equal-length shards - init_method = get_distributed_init_method() + init_method = get_file_store_init_method() with tempfile.NamedTemporaryFile(delete=False, suffix=".npz") as f_in: input_file = f_in.name @@ -1036,8 +1036,8 @@ def test_ulysses_uaa_2d_mask_layout_matches_baseline( num_heads = 3 head_size = 8 - base_init_method = get_distributed_init_method() - sp_init_method = get_distributed_init_method() + base_init_method = get_file_store_init_method() + sp_init_method = get_file_store_init_method() with tempfile.NamedTemporaryFile(delete=False, suffix=".npz") as f_in: input_file = f_in.name @@ -1150,8 +1150,8 @@ def test_ulysses_uaa_hybrid_ring_matches_baseline( # rank0/1 -> 3+2=5, rank2/3 -> 3+2=5 split_sizes = [3, 2, 3, 2] - base_init_method = get_distributed_init_method() - sp_init_method = get_distributed_init_method() + base_init_method = get_file_store_init_method() + sp_init_method = get_file_store_init_method() with tempfile.NamedTemporaryFile(delete=False, suffix=".npz") as f_in: input_file = f_in.name diff --git a/tests/diffusion/cache/test_seacache.py b/tests/diffusion/cache/test_seacache.py index ea11537d1..5f88db5c6 100644 --- a/tests/diffusion/cache/test_seacache.py +++ b/tests/diffusion/cache/test_seacache.py @@ -9,8 +9,8 @@ from types import SimpleNamespace import pytest import torch +from vllm.utils.network_utils import get_file_store_init_method -from tests.helpers.runtime import get_distributed_init_method from vllm_omni.diffusion.cache.seacache import ( SeaCacheBackend, SeaCacheConfig, @@ -375,7 +375,7 @@ def test_parameter_sharded_hook_handles_uneven_cfg_branch_dispatch() -> None: result_queue = manager.Queue() torch.multiprocessing.spawn( _uneven_cfg_sharded_worker, - args=(get_distributed_init_method("seacache_uneven_cfg_"), result_queue), + args=(get_file_store_init_method(), result_queue), nprocs=4, ) @@ -452,7 +452,7 @@ def test_hook_synchronizes_full_hybrid_sequence_parallel_group() -> None: result_queue = manager.Queue() torch.multiprocessing.spawn( _hybrid_sp_worker, - args=(get_distributed_init_method("seacache_hybrid_sp_"), result_queue), + args=(get_file_store_init_method(), result_queue), nprocs=4, ) diff --git a/tests/diffusion/distributed/test_hsdp.py b/tests/diffusion/distributed/test_hsdp.py index 433dec02f..3c3805a82 100644 --- a/tests/diffusion/distributed/test_hsdp.py +++ b/tests/diffusion/distributed/test_hsdp.py @@ -9,8 +9,8 @@ import torch import torch.distributed as dist import torch.nn as nn from torch.distributed.tensor import DeviceMesh, DTensor +from vllm.utils.network_utils import get_file_store_init_method -from tests.helpers.runtime import get_distributed_init_method from vllm_omni.diffusion.data import DiffusionParallelConfig from vllm_omni.diffusion.distributed import hsdp as hsdp_module from vllm_omni.diffusion.distributed.hsdp import ( @@ -43,7 +43,7 @@ def cpu_process_group(): yield return - dist.init_process_group("gloo", rank=0, world_size=1, init_method=get_distributed_init_method()) + dist.init_process_group("gloo", rank=0, world_size=1, init_method=get_file_store_init_method()) try: yield finally: diff --git a/tests/diffusion/distributed/test_parallel_state_sp_groups.py b/tests/diffusion/distributed/test_parallel_state_sp_groups.py index b42e97ba2..3aaefbcf9 100644 --- a/tests/diffusion/distributed/test_parallel_state_sp_groups.py +++ b/tests/diffusion/distributed/test_parallel_state_sp_groups.py @@ -9,8 +9,8 @@ from types import SimpleNamespace import pytest import torch import vllm.distributed.parallel_state as vllm_parallel_state +from vllm.utils.network_utils import get_file_store_init_method -from tests.helpers.runtime import get_distributed_init_method from vllm_omni.diffusion.distributed import parallel_state as omni_parallel_state from vllm_omni.diffusion.distributed.parallel_state import RankGenerator, set_seq_parallel_pg @@ -47,7 +47,7 @@ def test_omni_manages_vllm_distributed_state(monkeypatch): world_size=1, rank=0, local_rank=0, - distributed_init_method=get_distributed_init_method(), + distributed_init_method=get_file_store_init_method(), backend="gloo", ) try: diff --git a/tests/diffusion/distributed/test_pipeline_parallel.py b/tests/diffusion/distributed/test_pipeline_parallel.py index 08c3cc775..7cebbcc53 100644 --- a/tests/diffusion/distributed/test_pipeline_parallel.py +++ b/tests/diffusion/distributed/test_pipeline_parallel.py @@ -15,11 +15,11 @@ import pytest import torch from vllm.model_executor.models.utils import PPMissingLayer, make_empty_intermediate_tensors_factory, make_layers from vllm.sequence import IntermediateTensors +from vllm.utils.network_utils import get_file_store_init_method from vllm.v1.worker.gpu_worker import AsyncIntermediateTensors import vllm_omni.diffusion.distributed.pipeline_parallel as pp_module from tests.helpers.mark import hardware_marks -from tests.helpers.runtime import get_distributed_init_method from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin from vllm_omni.diffusion.distributed.parallel_state import ( destroy_distributed_env, @@ -520,7 +520,7 @@ def _run_isend_irecv(pp_size: int, device_kind: DeviceKind, init_method: str) -> @pytest.mark.parametrize("pp_size", [2]) def test_isend_irecv_tensor_dict(pp_size: int): """isend_tensor_dict / irecv_tensor_dict transfer a tensor dict without loss.""" - _run_isend_irecv(pp_size, device_kind="cpu", init_method=get_distributed_init_method()) + _run_isend_irecv(pp_size, device_kind="cpu", init_method=get_file_store_init_method()) @pytest.mark.full_model @@ -529,7 +529,7 @@ def test_isend_irecv_tensor_dict(pp_size: int): @pytest.mark.parametrize("pp_size", [pytest.param(2, marks=_L4_TWO_GPU)]) def test_isend_irecv_tensor_dict_parity(pp_size: int): """Nightly: isend/irecv on real multi-GPU NCCL.""" - _run_isend_irecv(pp_size, device_kind="cuda", init_method=get_distributed_init_method()) + _run_isend_irecv(pp_size, device_kind="cuda", init_method=get_file_store_init_method()) # --------------------------------------------------------------------------- @@ -560,7 +560,7 @@ def compute_single_gpu_baseline( if key in _baseline_cache: return _baseline_cache[key] - device = init_dist(0, 1, get_distributed_init_method(), device_kind) + device = init_dist(0, 1, get_file_store_init_method(), device_kind) try: initialize_model_parallel(pipeline_parallel_size=1, backend=_mp_backend(device_kind)) @@ -768,7 +768,7 @@ def test_predict_noise(pp_size, cfg_size, do_true_cfg, num_layers, input_seed, r rtol=rtol, atol=atol, device_kind="cpu", - init_method=get_distributed_init_method(), + init_method=get_file_store_init_method(), ) @@ -791,7 +791,7 @@ def test_predict_noise_parity(pp_size, cfg_size, do_true_cfg, dtype, num_layers, rtol=rtol, atol=atol, device_kind="cuda", - init_method=get_distributed_init_method(), + init_method=get_file_store_init_method(), ) @@ -806,7 +806,7 @@ def compute_scheduler_step_baseline( device_kind: DeviceKind, ) -> torch.Tensor: """Single-process reference: predict_noise + scheduler_step.""" - device = init_dist(0, 1, get_distributed_init_method(), device_kind) + device = init_dist(0, 1, get_file_store_init_method(), device_kind) try: initialize_model_parallel(pipeline_parallel_size=1, backend=_mp_backend(device_kind)) @@ -939,7 +939,7 @@ def test_scheduler_step(pp_size, cfg_size, do_true_cfg, input_seed): do_true_cfg=do_true_cfg, input_seed=input_seed, device_kind="cpu", - init_method=get_distributed_init_method(), + init_method=get_file_store_init_method(), ) @@ -961,5 +961,5 @@ def test_scheduler_step_parity(pp_size, cfg_size, do_true_cfg, input_seed): do_true_cfg=do_true_cfg, input_seed=input_seed, device_kind="cuda", - init_method=get_distributed_init_method(), + init_method=get_file_store_init_method(), ) diff --git a/tests/diffusion/distributed/test_ulysses_uaa_perf.py b/tests/diffusion/distributed/test_ulysses_uaa_perf.py index 707800cc0..621e59e43 100644 --- a/tests/diffusion/distributed/test_ulysses_uaa_perf.py +++ b/tests/diffusion/distributed/test_ulysses_uaa_perf.py @@ -14,9 +14,9 @@ from dataclasses import dataclass import pytest import torch import torch.distributed as dist +from vllm.utils.network_utils import get_file_store_init_method from tests.helpers.mark import hardware_test -from tests.helpers.runtime import get_distributed_init_method from vllm_omni.diffusion.attention.parallel.ulysses import ( _all_gather_int, _ulysses_all_to_all_any_o, @@ -62,7 +62,7 @@ def test_ulysses_advanced_uaa_comm_overhead(case: _PerfCase) -> None: if available_gpus < case.world_size: pytest.skip(f"Requires {case.world_size} GPUs, got {available_gpus}") - init_method = get_distributed_init_method() + init_method = get_file_store_init_method() torch.multiprocessing.spawn( _perf_worker, args=(case.world_size, init_method, case.ulysses_degree, case.ring_degree), diff --git a/tests/diffusion/offloader/conftest.py b/tests/diffusion/offloader/conftest.py index 6a55b7284..42c8841b7 100644 --- a/tests/diffusion/offloader/conftest.py +++ b/tests/diffusion/offloader/conftest.py @@ -7,14 +7,14 @@ import gc import pytest import torch.distributed as dist +from vllm.utils.network_utils import get_file_store_init_method -from tests.helpers.runtime import get_distributed_init_method from vllm_omni.platforms import current_omni_platform @pytest.fixture(scope="module") def dist_group(): - dist.init_process_group("gloo", rank=0, world_size=1, init_method=get_distributed_init_method()) + dist.init_process_group("gloo", rank=0, world_size=1, init_method=get_file_store_init_method()) try: yield finally: diff --git a/tests/helpers/runtime.py b/tests/helpers/runtime.py index 2a6addd93..5fedd17dd 100644 --- a/tests/helpers/runtime.py +++ b/tests/helpers/runtime.py @@ -78,20 +78,6 @@ def get_open_port(host: str = "127.0.0.1", *, max_attempts: int = 128) -> int: ) from last_exc -def get_distributed_init_method(prefix: str = "torch_dist_init_") -> str: - """Return a ``file://`` init_method for a ``torch.distributed`` process group. - - Args: - prefix: Prefix for the temporary rendezvous filename. Defaults to - ``torch_dist_init_``. - - Returns: - An init_method string for ``torch.distributed.init_process_group``. - """ - with tempfile.NamedTemporaryFile(prefix=prefix) as f: - return f"file://{f.name}" - - def dummy_messages_from_mix_data( system_prompt: dict[str, Any] | None = None, video_data_url: Any = None,