anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

[BugFix] Eliminate distributed test port race by switching to file:// rendezvous (#6468)

Browse Source 
Signed-off-by: Nick Cao <ncao@redhat.com>
This commit is contained in:
Nick Cao authored and GitHub committed 2026-08-23 07:08:44 +08:00
1 parent 3ad64c6a12
commit b2c9376539
11 files changed
+117 -237
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

tests/diffusion/distributed/test_pipeline_parallel.py
+31 -49
View File
@@ -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(),
)