[Tests] Use get_file_store_init_method helper from vllm (#7528)

Signed-off-by: Nick Cao <ncao@redhat.com>
Co-authored-by: Codex <noreply@openai.com>

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

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,
     )
 

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:

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:

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(),
     )

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),

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:

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,