diff --git a/tests/diffusion/attention/test_ulysses_uaa.py b/tests/diffusion/attention/test_ulysses_uaa.py index 184807d98..77e00edda 100644 --- a/tests/diffusion/attention/test_ulysses_uaa.py +++ b/tests/diffusion/attention/test_ulysses_uaa.py @@ -1,10 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from __future__ import annotations import os -import socket import tempfile from types import SimpleNamespace @@ -12,6 +11,7 @@ import numpy as np import pytest import torch +from tests.helpers.runtime import get_distributed_init_method from vllm_omni.diffusion.attention.layer import Attention from vllm_omni.diffusion.attention.parallel.ulysses import UlyssesParallelAttention from vllm_omni.diffusion.config import set_current_diffusion_config @@ -27,20 +27,6 @@ from vllm_omni.platforms import current_omni_platform pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] -def _find_free_port() -> int: - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: - s.bind(("127.0.0.1", 0)) - return int(s.getsockname()[1]) - - -def _set_dist_env(*, rank: int, world_size: int, master_port: int) -> None: - os.environ["RANK"] = str(rank) - os.environ["LOCAL_RANK"] = str(rank) - os.environ["WORLD_SIZE"] = str(world_size) - os.environ["MASTER_ADDR"] = "127.0.0.1" - os.environ["MASTER_PORT"] = str(master_port) - - @pytest.mark.core_model @pytest.mark.cpu def test_advanced_uaa_hybrid_rejects_padded_mqa_heads(monkeypatch) -> None: @@ -71,7 +57,7 @@ def test_advanced_uaa_hybrid_rejects_padded_mqa_heads(monkeypatch) -> None: def _run_attention_case( local_rank: int, world_size: int, - master_port: int, + init_method: str, input_file: str, output_file: str, num_heads: int, @@ -85,10 +71,9 @@ def _run_attention_case( device = torch.device(f"{current_omni_platform.device_type}:{local_rank}") current_omni_platform.set_device(device) - _set_dist_env(rank=local_rank, world_size=world_size, master_port=master_port) os.environ["DIFFUSION_ATTENTION_BACKEND"] = "TORCH_SDPA" - init_distributed_environment(world_size=world_size, rank=local_rank) + init_distributed_environment(world_size=world_size, rank=local_rank, distributed_init_method=init_method) initialize_model_parallel( data_parallel_size=1, cfg_parallel_size=1, @@ -202,10 +187,8 @@ def test_ulysses_uaa_matches_baseline(sp_world_size: int, seq_len: int, num_head batch_size = 2 head_size = 8 - base_port = _find_free_port() - sp_port = _find_free_port() - while sp_port == base_port: - sp_port = _find_free_port() + base_init_method = get_distributed_init_method() + sp_init_method = get_distributed_init_method() with tempfile.NamedTemporaryFile(delete=False, suffix=".npz") as f_in: input_file = f_in.name @@ -224,7 +207,7 @@ def test_ulysses_uaa_matches_baseline(sp_world_size: int, seq_len: int, num_head # Baseline (no SP) torch.multiprocessing.spawn( _run_attention_case, - args=(1, base_port, input_file, baseline_file, num_heads, head_size, 1, "strict"), + args=(1, base_init_method, input_file, baseline_file, num_heads, head_size, 1, "strict"), nprocs=1, ) @@ -233,7 +216,7 @@ def test_ulysses_uaa_matches_baseline(sp_world_size: int, seq_len: int, num_head _run_attention_case, args=( sp_world_size, - sp_port, + sp_init_method, input_file, sp_file, num_heads, @@ -278,10 +261,8 @@ def test_ulysses_uaa_hybrid_ring_matches_baseline() -> None: # rank0/1 -> 3+2=5, rank2/3 -> 3+2=5 split_sizes = [3, 2, 3, 2] - base_port = _find_free_port() - sp_port = _find_free_port() - while sp_port == base_port: - sp_port = _find_free_port() + base_init_method = get_distributed_init_method() + sp_init_method = get_distributed_init_method() with tempfile.NamedTemporaryFile(delete=False, suffix=".npz") as f_in: input_file = f_in.name @@ -300,7 +281,19 @@ def test_ulysses_uaa_hybrid_ring_matches_baseline() -> None: # Baseline (no SP) torch.multiprocessing.spawn( _run_attention_case, - args=(1, base_port, input_file, baseline_file, num_heads, head_size, 1, "strict", 1, None, "mem_efficient"), + args=( + 1, + base_init_method, + input_file, + baseline_file, + num_heads, + head_size, + 1, + "strict", + 1, + None, + "mem_efficient", + ), nprocs=1, ) @@ -309,7 +302,7 @@ def test_ulysses_uaa_hybrid_ring_matches_baseline() -> None: _run_attention_case, args=( sp_world_size, - sp_port, + sp_init_method, input_file, sp_file, num_heads, diff --git a/tests/diffusion/distributed/test_hsdp.py b/tests/diffusion/distributed/test_hsdp.py index 0703d19b8..25762c356 100644 --- a/tests/diffusion/distributed/test_hsdp.py +++ b/tests/diffusion/distributed/test_hsdp.py @@ -1,10 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Unit tests for HSDP (Hybrid Sharded Data Parallel) configuration and utilities.""" import gc -import os -import socket import pytest import torch @@ -12,6 +10,7 @@ import torch.distributed as dist import torch.nn as nn from torch.distributed.tensor import DeviceMesh, DTensor +from tests.helpers.runtime import get_distributed_init_method from vllm_omni.diffusion.data import DiffusionParallelConfig from vllm_omni.diffusion.distributed.hsdp import ( HSDPInferenceConfig, @@ -37,35 +36,17 @@ class _PackedModel(nn.Module): self.root_weight = nn.Parameter(torch.ones(2)) -def _find_free_port() -> int: - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: - sock.bind(("127.0.0.1", 0)) - return int(sock.getsockname()[1]) - - @pytest.fixture(scope="module") def cpu_process_group(): if dist.is_initialized(): yield return - master_port = _find_free_port() - os.environ.update( - { - "RANK": "0", - "LOCAL_RANK": "0", - "WORLD_SIZE": "1", - "MASTER_ADDR": "127.0.0.1", - "MASTER_PORT": str(master_port), - } - ) - dist.init_process_group("gloo", rank=0, world_size=1) + dist.init_process_group("gloo", rank=0, world_size=1, init_method=get_distributed_init_method()) try: yield finally: dist.destroy_process_group() - for key in ("MASTER_ADDR", "MASTER_PORT", "RANK", "WORLD_SIZE", "LOCAL_RANK"): - os.environ.pop(key, None) gc.collect() diff --git a/tests/diffusion/distributed/test_pipeline_parallel.py b/tests/diffusion/distributed/test_pipeline_parallel.py index 0b48ed266..cccf116c3 100644 --- a/tests/diffusion/distributed/test_pipeline_parallel.py +++ b/tests/diffusion/distributed/test_pipeline_parallel.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Tests for PipelineParallelMixin and related PP communication helpers. CPU unit tests exercise mixin logic without a process group. Distributed CPU @@ -9,8 +9,6 @@ real multi-GPU NCCL collectives. from __future__ import annotations -import os -import socket from typing import Literal import pytest @@ -21,6 +19,7 @@ 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, @@ -40,31 +39,18 @@ DeviceKind = Literal["cpu", "cuda"] _UNIT_MARKS = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] -def _find_free_port() -> str: - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: - s.bind(("127.0.0.1", 0)) - return str(s.getsockname()[1]) - - -def update_environment_variables(envs_dict: dict[str, str]) -> None: - for k, v in envs_dict.items(): - os.environ[k] = v - - def _cleanup_distributed() -> None: - """Destroy omni distributed state and clear test-process leftovers. + """Destroy omni distributed state left dangling after a test. ``destroy_distributed_env`` leaves ``vllm.distributed.parallel_state._PP`` aliased to the destroyed group. Parent-process baselines must clear that - dangling reference (and RANK/MASTER_* env) so later unit tests that call + dangling reference so later unit tests that call ``vllm.distributed.initialize_model_parallel`` do not see stale state. """ destroy_distributed_env() import vllm.distributed.parallel_state as vllm_parallel_state vllm_parallel_state._PP = None - for key in ("MASTER_ADDR", "MASTER_PORT", "RANK", "WORLD_SIZE", "LOCAL_RANK"): - os.environ.pop(key, None) # --------------------------------------------------------------------------- @@ -416,24 +402,20 @@ def _worker_device(local_rank: int, device_kind: DeviceKind) -> torch.device: def init_dist( local_rank: int, world_size: int, - master_port: str, + init_method: str, device_kind: DeviceKind, ) -> torch.device: """Initialise the distributed environment for a spawned worker.""" device = _worker_device(local_rank, device_kind) if device_kind == "cuda": current_omni_platform.set_device(device) - update_environment_variables( - { - "RANK": str(local_rank), - "LOCAL_RANK": str(local_rank), - "WORLD_SIZE": str(world_size), - "MASTER_ADDR": "localhost", - "MASTER_PORT": master_port, - } - ) backend = "gloo" if device_kind == "cpu" else None - init_distributed_environment(backend=backend) + init_distributed_environment( + world_size=world_size, + rank=local_rank, + distributed_init_method=init_method, + backend=backend, + ) return device @@ -493,11 +475,11 @@ def make_pipeline_and_inputs( def isend_irecv_worker( local_rank: int, world_size: int, - master_port: str, + init_method: str, device_kind: DeviceKind, result_queue, ): - device = init_dist(local_rank, world_size, master_port, device_kind) + device = init_dist(local_rank, world_size, init_method, device_kind) initialize_model_parallel(pipeline_parallel_size=world_size, backend=_mp_backend(device_kind)) pp_group = get_pp_group() @@ -517,13 +499,13 @@ def isend_irecv_worker( _cleanup_distributed() -def _run_isend_irecv(pp_size: int, device_kind: DeviceKind, master_port: str) -> None: +def _run_isend_irecv(pp_size: int, device_kind: DeviceKind, init_method: str) -> None: mp_context = torch.multiprocessing.get_context("spawn") manager = mp_context.Manager() q = manager.Queue() torch.multiprocessing.spawn( isend_irecv_worker, - args=(pp_size, master_port, device_kind, q), + args=(pp_size, init_method, device_kind, q), nprocs=pp_size, ) results = {label: tensor for label, tensor in [q.get(), q.get()]} @@ -538,7 +520,7 @@ def _run_isend_irecv(pp_size: int, device_kind: DeviceKind, master_port: 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", master_port=_find_free_port()) + _run_isend_irecv(pp_size, device_kind="cpu", init_method=get_distributed_init_method()) @pytest.mark.full_model @@ -547,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", master_port=_find_free_port()) + _run_isend_irecv(pp_size, device_kind="cuda", init_method=get_distributed_init_method()) # --------------------------------------------------------------------------- @@ -578,7 +560,7 @@ def compute_single_gpu_baseline( if key in _baseline_cache: return _baseline_cache[key] - device = init_dist(0, 1, _find_free_port(), device_kind) + device = init_dist(0, 1, get_distributed_init_method(), device_kind) try: initialize_model_parallel(pipeline_parallel_size=1, backend=_mp_backend(device_kind)) @@ -604,7 +586,7 @@ def compute_single_gpu_baseline( def predict_noise_worker( local_rank: int, world_size: int, - master_port: str, + init_method: str, pp_size: int, cfg_size: int, do_true_cfg: bool, @@ -614,7 +596,7 @@ def predict_noise_worker( result_queue, ): """Generic predict-noise worker parameterized by PP and CFG topology.""" - device = init_dist(local_rank, world_size, master_port, device_kind) + device = init_dist(local_rank, world_size, init_method, device_kind) initialize_model_parallel( pipeline_parallel_size=pp_size, cfg_parallel_size=cfg_size, @@ -661,7 +643,7 @@ def _run_predict_noise( rtol: float, atol: float, device_kind: DeviceKind, - master_port: str, + init_method: str, ) -> None: test_config = { "num_layers": num_layers, @@ -681,7 +663,7 @@ def _run_predict_noise( world_size = pp_size * cfg_size torch.multiprocessing.spawn( predict_noise_worker, - args=(world_size, master_port, pp_size, cfg_size, do_true_cfg, dtype, test_config, device_kind, pp_q), + args=(world_size, init_method, pp_size, cfg_size, do_true_cfg, dtype, test_config, device_kind, pp_q), nprocs=world_size, ) @@ -786,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", - master_port=_find_free_port(), + init_method=get_distributed_init_method(), ) @@ -809,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", - master_port=_find_free_port(), + init_method=get_distributed_init_method(), ) @@ -824,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, _find_free_port(), device_kind) + device = init_dist(0, 1, get_distributed_init_method(), device_kind) try: initialize_model_parallel(pipeline_parallel_size=1, backend=_mp_backend(device_kind)) @@ -853,7 +835,7 @@ def compute_scheduler_step_baseline( def scheduler_step_worker( local_rank: int, world_size: int, - master_port: str, + init_method: str, pp_size: int, cfg_size: int, do_true_cfg: bool, @@ -861,7 +843,7 @@ def scheduler_step_worker( device_kind: DeviceKind, result_queue, ): - device = init_dist(local_rank, world_size, master_port, device_kind) + device = init_dist(local_rank, world_size, init_method, device_kind) initialize_model_parallel( pipeline_parallel_size=pp_size, cfg_parallel_size=cfg_size, @@ -905,7 +887,7 @@ def _run_scheduler_step( do_true_cfg: bool, input_seed: int, device_kind: DeviceKind, - master_port: str, + init_method: str, ) -> None: test_config = { "num_layers": 4, @@ -925,7 +907,7 @@ def _run_scheduler_step( world_size = pp_size * cfg_size torch.multiprocessing.spawn( scheduler_step_worker, - args=(world_size, master_port, pp_size, cfg_size, do_true_cfg, test_config, device_kind, q), + args=(world_size, init_method, pp_size, cfg_size, do_true_cfg, test_config, device_kind, q), nprocs=world_size, ) @@ -957,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", - master_port=_find_free_port(), + init_method=get_distributed_init_method(), ) @@ -979,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", - master_port=_find_free_port(), + init_method=get_distributed_init_method(), ) diff --git a/tests/diffusion/distributed/test_ulysses_uaa_perf.py b/tests/diffusion/distributed/test_ulysses_uaa_perf.py index 2a16a9ae5..707800cc0 100644 --- a/tests/diffusion/distributed/test_ulysses_uaa_perf.py +++ b/tests/diffusion/distributed/test_ulysses_uaa_perf.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Perf smoke tests for Ulysses advanced_uaa communication overhead. This test is intended for CI monitoring only: @@ -9,8 +9,6 @@ This test is intended for CI monitoring only: from __future__ import annotations -import os -import socket from dataclasses import dataclass import pytest @@ -18,6 +16,7 @@ import torch import torch.distributed as dist 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, @@ -33,20 +32,6 @@ from vllm_omni.diffusion.distributed.parallel_state import ( from vllm_omni.platforms import current_omni_platform -def _find_free_port() -> int: - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: - s.bind(("127.0.0.1", 0)) - return int(s.getsockname()[1]) - - -def _set_dist_env(*, rank: int, world_size: int, master_port: int) -> None: - os.environ["RANK"] = str(rank) - os.environ["LOCAL_RANK"] = str(rank) - os.environ["WORLD_SIZE"] = str(world_size) - os.environ["MASTER_ADDR"] = "127.0.0.1" - os.environ["MASTER_PORT"] = str(master_port) - - def _max_all_reduce(pg: dist.ProcessGroup, value: float, *, device: torch.device) -> float: t = torch.tensor([value], device=device, dtype=torch.float32) dist.all_reduce(t, op=dist.ReduceOp.MAX, group=pg) @@ -77,22 +62,20 @@ 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}") - master_port = _find_free_port() + init_method = get_distributed_init_method() torch.multiprocessing.spawn( _perf_worker, - args=(case.world_size, master_port, case.ulysses_degree, case.ring_degree), + args=(case.world_size, init_method, case.ulysses_degree, case.ring_degree), nprocs=case.world_size, ) -def _perf_worker(local_rank: int, world_size: int, master_port: int, ulysses_degree: int, ring_degree: int) -> None: +def _perf_worker(local_rank: int, world_size: int, init_method: str, ulysses_degree: int, ring_degree: int) -> None: device = torch.device(f"{current_omni_platform.device_type}:{local_rank}") current_omni_platform.set_device(device) - _set_dist_env(rank=local_rank, world_size=world_size, master_port=master_port) - try: - init_distributed_environment(world_size=world_size, rank=local_rank) + init_distributed_environment(world_size=world_size, rank=local_rank, distributed_init_method=init_method) initialize_model_parallel( data_parallel_size=1, cfg_parallel_size=1, diff --git a/tests/diffusion/offloader/test_distributed_layerwise_backend.py b/tests/diffusion/offloader/test_distributed_layerwise_backend.py index c67a2d45b..78dd82b6f 100644 --- a/tests/diffusion/offloader/test_distributed_layerwise_backend.py +++ b/tests/diffusion/offloader/test_distributed_layerwise_backend.py @@ -1,12 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Unit tests for DistributedLayerwiseOffloadHook and backend utilities.""" import gc import json -import os -import socket from contextlib import contextmanager from types import SimpleNamespace @@ -18,6 +16,7 @@ from torch import nn from torch.distributed.tensor import DeviceMesh, DTensor, Replicate import vllm_omni.diffusion.offloader.distributed_layerwise_backend as dist_backend_module +from tests.helpers.runtime import get_distributed_init_method from vllm_omni.diffusion.model_loader.host_weight_plan import ( HostWeightPlan, TensorBinding, @@ -65,27 +64,10 @@ def dummy_stream(_stream): yield None -def _find_free_port() -> int: - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: - s.bind(("127.0.0.1", 0)) - return int(s.getsockname()[1]) - - -def _set_dist_env(*, rank: int, world_size: int, master_port: int) -> None: - os.environ["RANK"] = str(rank) - os.environ["LOCAL_RANK"] = str(rank) - os.environ["WORLD_SIZE"] = str(world_size) - os.environ["MASTER_ADDR"] = "127.0.0.1" - os.environ["MASTER_PORT"] = str(master_port) - - def _cleanup_distributed() -> None: if dist.is_initialized(): dist.destroy_process_group() - for key in ["MASTER_ADDR", "MASTER_PORT", "RANK", "WORLD_SIZE", "LOCAL_RANK"]: - os.environ.pop(key, None) - gc.collect() if current_omni_platform.is_available(): current_omni_platform.empty_cache() @@ -94,10 +76,7 @@ def _cleanup_distributed() -> None: @pytest.fixture(scope="module") def dist_group(): - master_port = _find_free_port() - _set_dist_env(rank=0, world_size=1, master_port=master_port) - - dist.init_process_group("gloo", rank=0, world_size=1) + dist.init_process_group("gloo", rank=0, world_size=1, init_method=get_distributed_init_method()) try: yield finally: diff --git a/tests/diffusion/offloader/test_layerwise_backend.py b/tests/diffusion/offloader/test_layerwise_backend.py index f62bd6472..9253f6d03 100644 --- a/tests/diffusion/offloader/test_layerwise_backend.py +++ b/tests/diffusion/offloader/test_layerwise_backend.py @@ -1,11 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Unit tests for LayerwiseOffloadHook and LayerWiseOffloadBackend utilities.""" import gc -import os -import socket from contextlib import contextmanager import pytest @@ -15,6 +13,7 @@ from torch import nn from torch.distributed.tensor import DeviceMesh, DTensor, Replicate import vllm_omni.diffusion.offloader.layerwise_backend as layerwise_backend_module +from tests.helpers.runtime import get_distributed_init_method from vllm_omni.diffusion.offloader.layerwise_backend import LayerWiseOffloadBackend, LayerwiseOffloadHook from vllm_omni.platforms import current_omni_platform @@ -39,27 +38,10 @@ def dummy_stream(_stream): yield None -def _find_free_port() -> int: - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: - s.bind(("127.0.0.1", 0)) - return int(s.getsockname()[1]) - - -def _set_dist_env(*, rank: int, world_size: int, master_port: int) -> None: - os.environ["RANK"] = str(rank) - os.environ["LOCAL_RANK"] = str(rank) - os.environ["WORLD_SIZE"] = str(world_size) - os.environ["MASTER_ADDR"] = "127.0.0.1" - os.environ["MASTER_PORT"] = str(master_port) - - def _cleanup_distributed() -> None: if dist.is_initialized(): dist.destroy_process_group() - for key in ["MASTER_ADDR", "MASTER_PORT", "RANK", "WORLD_SIZE", "LOCAL_RANK"]: - os.environ.pop(key, None) - gc.collect() if current_omni_platform.is_available(): current_omni_platform.empty_cache() @@ -68,10 +50,7 @@ def _cleanup_distributed() -> None: @pytest.fixture(scope="module") def dist_group(): - master_port = _find_free_port() - _set_dist_env(rank=0, world_size=1, master_port=master_port) - - dist.init_process_group("gloo", rank=0, world_size=1) + dist.init_process_group("gloo", rank=0, world_size=1, init_method=get_distributed_init_method()) try: yield finally: diff --git a/tests/diffusion/test_omni_diffusion_config_master_port.py b/tests/diffusion/test_omni_diffusion_config_master_port.py index 8addd1c8c..0fe2dd08d 100644 --- a/tests/diffusion/test_omni_diffusion_config_master_port.py +++ b/tests/diffusion/test_omni_diffusion_config_master_port.py @@ -1,22 +1,17 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Unit tests for OmniDiffusionConfig master_port resolution (issue #3794).""" import socket import pytest +from tests.helpers.runtime import get_open_port from vllm_omni.diffusion.data import OmniDiffusionConfig pytestmark = [pytest.mark.core_model, pytest.mark.cpu] -def _get_available_port() -> int: - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: - sock.bind(("", 0)) - return sock.getsockname()[1] - - def _bind_available_port_below(max_port: int, port_inc: int) -> socket.socket: for _ in range(100): sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) @@ -36,27 +31,27 @@ def _bind_available_port_below(max_port: int, port_inc: int) -> socket.socket: class TestOmniDiffusionConfigMasterPort: def test_honors_master_port_env(self, monkeypatch: pytest.MonkeyPatch) -> None: - port = _get_available_port() + port = get_open_port() monkeypatch.setenv("MASTER_PORT", str(port)) config = OmniDiffusionConfig(model="test") assert config.master_port == port def test_honors_explicit_master_port(self, monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.delenv("MASTER_PORT", raising=False) - port = _get_available_port() + port = get_open_port() config = OmniDiffusionConfig(model="test", master_port=port) assert config.master_port == port def test_master_port_env_takes_precedence_over_kwarg(self, monkeypatch: pytest.MonkeyPatch) -> None: - env_port = _get_available_port() - kwarg_port = _get_available_port() + env_port = get_open_port() + kwarg_port = get_open_port() monkeypatch.setenv("MASTER_PORT", str(env_port)) config = OmniDiffusionConfig(model="test", master_port=kwarg_port) assert config.master_port == env_port def test_explicit_master_port_is_stable_without_jitter(self, monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.delenv("MASTER_PORT", raising=False) - port = _get_available_port() + port = get_open_port() ports = {OmniDiffusionConfig(model="test", master_port=port).master_port for _ in range(5)} assert ports == {port} diff --git a/tests/distributed/omni_connectors/test_bagel_mooncake_connector.py b/tests/distributed/omni_connectors/test_bagel_mooncake_connector.py index 3b9fcbb64..c31b80a6f 100644 --- a/tests/distributed/omni_connectors/test_bagel_mooncake_connector.py +++ b/tests/distributed/omni_connectors/test_bagel_mooncake_connector.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ End-to-end test for Bagel text2img with the Mooncake inter-stage connector. @@ -22,7 +22,7 @@ import pytest from PIL import Image from tests.helpers.mark import hardware_test -from tests.helpers.runtime import OmniRunner +from tests.helpers.runtime import OmniRunner, get_open_port from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config from vllm_omni.entrypoints.omni import Omni from vllm_omni.outputs import OmniRequestOutput @@ -55,15 +55,6 @@ PIXEL_TOLERANCE = 10 DEFAULT_PROMPT = "<|im_start|>A cute cat<|im_end|>" -def _find_free_port() -> int: - """Find and return a free ephemeral port by binding to port 0.""" - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: - s.bind(("127.0.0.1", 0)) - s.listen(1) - port = s.getsockname()[1] - return port - - def _configure_sampling_params(omni: Omni, num_inference_steps: int = 15) -> list: """Configure sampling parameters for Bagel text2img generation. @@ -277,9 +268,9 @@ def test_bagel_text2img_mooncake_connector(run_level): "mooncake_master is not available or cannot execute (missing shared libraries like libibverbs)" ) MOONCAKE_HOST = "127.0.0.1" - MOONCAKE_RPC_PORT = _find_free_port() - MOONCAKE_HTTP_PORT = _find_free_port() - MOONCAKE_METRICS_PORT = _find_free_port() + MOONCAKE_RPC_PORT = get_open_port() + MOONCAKE_HTTP_PORT = get_open_port() + MOONCAKE_METRICS_PORT = get_open_port() mooncake_master_proc = None temp_config_file = None diff --git a/tests/distributed/omni_connectors/test_bagel_shared_memory_connector.py b/tests/distributed/omni_connectors/test_bagel_shared_memory_connector.py index 72b56d3e9..6b23ba321 100644 --- a/tests/distributed/omni_connectors/test_bagel_shared_memory_connector.py +++ b/tests/distributed/omni_connectors/test_bagel_shared_memory_connector.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ End-to-end tests for Bagel with shared memory connector: img2img and text2img. @@ -13,7 +13,6 @@ End-to-end tests for Bagel with shared memory connector: img2img and text2img. import os os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" -import socket from typing import Any import pytest @@ -82,15 +81,6 @@ def _load_input_image() -> Image.Image: return ImageAsset("2560px-Gfp-wisconsin-madison-the-nature-boardwalk").pil_image.convert("RGB") -def _find_free_port() -> int: - """Find and return a free ephemeral port by binding to port 0.""" - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: - s.bind(("127.0.0.1", 0)) - s.listen(1) - port = s.getsockname()[1] - return port - - def _configure_sampling_params(omni: Omni, num_inference_steps: int = 15) -> list: """Configure sampling parameters for Bagel img2img generation. diff --git a/tests/distributed/omni_connectors/test_mooncake_transfer_engine_rdma.py b/tests/distributed/omni_connectors/test_mooncake_transfer_engine_rdma.py index 60653b3cd..66f0327f7 100755 --- a/tests/distributed/omni_connectors/test_mooncake_transfer_engine_rdma.py +++ b/tests/distributed/omni_connectors/test_mooncake_transfer_engine_rdma.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ Integration tests for MooncakeTransferEngineConnector. @@ -18,6 +18,7 @@ import pytest import torch from tests.helpers.mark import hardware_test +from tests.helpers.runtime import get_open_port from vllm_omni.distributed.omni_connectors.connectors.mooncake_transfer_engine_connector import ( ManagedBuffer, MooncakeTransferEngineConnector, @@ -164,14 +165,6 @@ RDMA_HOST = get_rdma_host() RDMA_DEVICE = _detect_rdma_device() -def _free_port() -> int: - import socket - - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: - s.bind((RDMA_HOST, 0)) - return s.getsockname()[1] - - def _connector_config( zmq_port: int, pool_size: int = 16 * 1024 * 1024, @@ -204,7 +197,7 @@ class TestBasicConnector: """Verify connector initialization, put, cleanup, and health check.""" def test_initialization(self): - port = _free_port() + port = get_open_port(RDMA_HOST) with MooncakeTransferEngineConnector(_connector_config(port, pool_size=1024 * 1024)) as c: assert c.rpc_port != 0 assert c.pool_size == 1024 * 1024 @@ -214,7 +207,7 @@ class TestBasicConnector: def test_put_tensor_bytes_object(self): """Put tensor / bytes / dict and verify metadata.""" - port = _free_port() + port = get_open_port(RDMA_HOST) with MooncakeTransferEngineConnector(_connector_config(port)) as c: ok, sz, meta = c.put("s0", "s1", "t", torch.randn(100)) assert ok @@ -229,7 +222,7 @@ class TestBasicConnector: assert not meta["is_fast_path"] def test_cleanup_releases_buffer(self): - port = _free_port() + port = get_open_port(RDMA_HOST) with MooncakeTransferEngineConnector(_connector_config(port)) as c: c.put("s0", "s1", "r1", torch.randn(100)) key = MooncakeTransferEngineConnector._make_key("r1", "s0", "s1") @@ -239,7 +232,7 @@ class TestBasicConnector: def test_pool_exhaustion_and_recovery(self): """Fill pool, verify failure, free, verify recovery.""" - port = _free_port() + port = get_open_port(RDMA_HOST) with MooncakeTransferEngineConnector(_connector_config(port, pool_size=64 * 1024)) as c: ids = [] for i in range(10): @@ -264,8 +257,8 @@ class TestEndToEnd: """E2E RDMA transfer: tensor, bytes, object, zero-copy, large payload, mixed types.""" def _pair(self, pool_size=16 * 1024 * 1024): - p = MooncakeTransferEngineConnector(_connector_config(_free_port(), pool_size)) - c = MooncakeTransferEngineConnector(_connector_config(_free_port(), pool_size)) + p = MooncakeTransferEngineConnector(_connector_config(get_open_port(RDMA_HOST), pool_size)) + c = MooncakeTransferEngineConnector(_connector_config(get_open_port(RDMA_HOST), pool_size)) return p, c def test_tensor_e2e(self): @@ -401,7 +394,7 @@ class TestEndToEnd: def test_concurrent_put(self): """10 concurrent puts should all succeed.""" - port = _free_port() + port = get_open_port(RDMA_HOST) conn = MooncakeTransferEngineConnector(_connector_config(port, pool_size=64 * 1024 * 1024)) errors: list[str] = [] lock = threading.Lock() @@ -457,20 +450,20 @@ class TestLifecycle: """Close, context manager, double-close safety.""" def test_close_releases_resources(self): - c = MooncakeTransferEngineConnector(_connector_config(_free_port(), pool_size=1024 * 1024)) + c = MooncakeTransferEngineConnector(_connector_config(get_open_port(RDMA_HOST), pool_size=1024 * 1024)) c.put("s0", "s1", "x", torch.randn(100)) c.close() assert c._stop_event.is_set() assert len(c._local_buffers) == 0 def test_context_manager(self): - with MooncakeTransferEngineConnector(_connector_config(_free_port())) as c: + with MooncakeTransferEngineConnector(_connector_config(get_open_port(RDMA_HOST))) as c: ok, _, _ = c.put("s0", "s1", "ctx", torch.randn(50)) assert ok assert c._stop_event.is_set() def test_double_close_safe(self): - c = MooncakeTransferEngineConnector(_connector_config(_free_port())) + c = MooncakeTransferEngineConnector(_connector_config(get_open_port(RDMA_HOST))) c.close() c.close() @@ -489,12 +482,12 @@ class TestGPUPool: return _connector_config(port, pool_size, pool_device="cuda:0") def test_gpu_pool_init(self): - with MooncakeTransferEngineConnector(self._gpu_cfg(_free_port())) as c: + with MooncakeTransferEngineConnector(self._gpu_cfg(get_open_port(RDMA_HOST))) as c: assert c.pool_device == "cuda:0" assert c.pool.is_cuda def test_gpu_pool_put_cpu_and_gpu_tensor(self): - with MooncakeTransferEngineConnector(self._gpu_cfg(_free_port())) as c: + with MooncakeTransferEngineConnector(self._gpu_cfg(get_open_port(RDMA_HOST))) as c: ok, _, meta = c.put("s0", "s1", "h2d", torch.randn(256, 256)) assert ok assert meta["is_fast_path"] @@ -504,8 +497,8 @@ class TestGPUPool: assert meta["is_fast_path"] def test_gpu_e2e_transfer(self): - p = MooncakeTransferEngineConnector(self._gpu_cfg(_free_port())) - c = MooncakeTransferEngineConnector(self._gpu_cfg(_free_port())) + p = MooncakeTransferEngineConnector(self._gpu_cfg(get_open_port(RDMA_HOST))) + c = MooncakeTransferEngineConnector(self._gpu_cfg(get_open_port(RDMA_HOST))) try: orig = torch.randn(512, 512, dtype=torch.float32, device="cuda:0") ok, _, meta = p.put("s0", "s1", "ge", orig) @@ -536,8 +529,8 @@ class TestStressCorrectness: """ def _pair(self, pool_size=64 * 1024 * 1024): - p = MooncakeTransferEngineConnector(_connector_config(_free_port(), pool_size)) - c = MooncakeTransferEngineConnector(_connector_config(_free_port(), pool_size)) + p = MooncakeTransferEngineConnector(_connector_config(get_open_port(RDMA_HOST), pool_size)) + c = MooncakeTransferEngineConnector(_connector_config(get_open_port(RDMA_HOST), pool_size)) return p, c # -- Concurrent put + get with data integrity -- @@ -712,7 +705,7 @@ class TestStressCorrectness: def test_empty_bytes_rejected(self): """Connector should gracefully reject empty bytes payload.""" - port = _free_port() + port = get_open_port(RDMA_HOST) with MooncakeTransferEngineConnector(_connector_config(port, pool_size=8 * 1024 * 1024)) as c: ok, sz, meta = c.put("s0", "s1", "empty_b", b"") assert not ok, "Empty bytes should be rejected by connector" @@ -744,7 +737,7 @@ class TestStressCorrectness: def test_rapid_alloc_free_cycle(self): """Put + cleanup in tight loop to stress allocator under real connector.""" - port = _free_port() + port = get_open_port(RDMA_HOST) with MooncakeTransferEngineConnector(_connector_config(port, pool_size=8 * 1024 * 1024)) as c: for i in range(50): rid = f"cycle_{i}" diff --git a/tests/helpers/runtime.py b/tests/helpers/runtime.py index 51df9ef84..d61dcc25f 100644 --- a/tests/helpers/runtime.py +++ b/tests/helpers/runtime.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """Server/client/runner runtime primitives for tests.""" import asyncio @@ -178,6 +181,17 @@ def get_open_port(host: str = "127.0.0.1", *, max_attempts: int = 128) -> int: ) from last_exc +def get_distributed_init_method() -> str: + """Return a fresh ``file://`` init_method for ``torch.distributed`` process groups. + + Rendezvous happens through a filesystem path instead of a pre-agreed TCP port, + so there's no bind-time race to retry around: the path is unique per call and + the real rank-0 process is the only one that ever creates it. + """ + with tempfile.NamedTemporaryFile(prefix="torch_dist_init_") as f: + return f"file://{f.name}" + + def dummy_messages_from_mix_data( system_prompt: dict[str, Any] = None, video_data_url: Any = None,