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