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/attention/test_ulysses_uaa.py
+24 -31
View File
@@ -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,
tests/diffusion/distributed/test_hsdp.py
+3 -22
View File
@@ -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()
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(),
)
tests/diffusion/distributed/test_ulysses_uaa_perf.py
+6 -23
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
"""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,
tests/diffusion/offloader/test_distributed_layerwise_backend.py
+3 -24
View File
@@ -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:
tests/diffusion/offloader/test_layerwise_backend.py
+3 -24
View File
@@ -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:
tests/diffusion/test_omni_diffusion_config_master_port.py
+7 -12
View File
@@ -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}
tests/distributed/omni_connectors/test_bagel_mooncake_connector.py
+5 -14
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
"""
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
tests/distributed/omni_connectors/test_bagel_shared_memory_connector.py
+1 -11
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
"""
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.
tests/distributed/omni_connectors/test_mooncake_transfer_engine_rdma.py
+20 -27
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
"""
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}"
tests/helpers/runtime.py
+14
View File
@@ -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,