[Core] Restore topology-aware diffusion worker titles (#6722)
Signed-off-by: hsliuustc0106 <hsliuustc0106@users.noreply.github.com> Co-authored-by: hsliuustc0106 <hsliuustc0106@users.noreply.github.com>
This commit is contained in:
Hongsheng Liuandhsliuustc0106 authored and GitHub committed
1 parent
6b9ce05feb
commit
8bfad23f98
4 files changed
+422
-9
No files matched your search
@@ -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
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
@@ -194,6 +194,47 @@ def test_moe_ep_maps_diffusion_sp_cfg_dp_to_vllm_groups(monkeypatch):
|
||||
assert ep_groups == [parallel_state.vllm_parallel_state._EP.local_group]
|
||||
|
||||
|
||||
@pytest.mark.cpu
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.parametrize(
|
||||
("fully_shard_degree", "error_message"),
|
||||
[(0, "fully_shard_degree must be positive"), (3, "must be divisible by fully_shard_degree")],
|
||||
)
|
||||
def test_invalid_hsdp_shard_size_fails_before_group_creation(monkeypatch, fully_shard_degree: int, error_message: str):
|
||||
"""Reject invalid HSDP shard sizes before allocating any process groups."""
|
||||
world_size = 4
|
||||
created_groups: list[object] = []
|
||||
|
||||
def fail_if_group_created(*args, **kwargs):
|
||||
del args, kwargs
|
||||
created_groups.append(object())
|
||||
raise AssertionError("a process group was created before HSDP validation")
|
||||
|
||||
fake_world_group = SimpleNamespace(
|
||||
rank_in_group=0,
|
||||
local_rank=0,
|
||||
device_group=object(),
|
||||
)
|
||||
monkeypatch.setattr(parallel_state.torch.distributed, "is_initialized", lambda: True)
|
||||
monkeypatch.setattr(parallel_state.torch.distributed, "get_world_size", lambda: world_size)
|
||||
monkeypatch.setattr(parallel_state, "get_world_group", lambda: fake_world_group)
|
||||
monkeypatch.setattr(parallel_state, "init_model_parallel_group", fail_if_group_created)
|
||||
|
||||
for name in ("_DP", "_CFG", "_SP", "_PP", "_FS", "_HSDP_REPLICATE", "_EXPERT_PARALLEL_GROUP_RANKS"):
|
||||
monkeypatch.setattr(parallel_state, name, None)
|
||||
for name in ("_TP", "_PCP", "_DP", "_EP", "_PP"):
|
||||
monkeypatch.setattr(parallel_state.vllm_parallel_state, name, None, raising=False)
|
||||
|
||||
with pytest.raises(ValueError, match=error_message):
|
||||
parallel_state.initialize_model_parallel(
|
||||
fully_shard_degree=fully_shard_degree,
|
||||
use_hsdp=True,
|
||||
backend="gloo",
|
||||
)
|
||||
|
||||
assert created_groups == []
|
||||
|
||||
|
||||
@pytest.mark.cpu
|
||||
@pytest.mark.core_model
|
||||
def test_cfg_parallel_keeps_diffusion_dp_without_ep(monkeypatch):
|
||||
|
||||
@@ -0,0 +1,255 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
import importlib.util
|
||||
import shutil
|
||||
import subprocess
|
||||
import sys
|
||||
import textwrap
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from unittest.mock import Mock
|
||||
|
||||
import pytest
|
||||
from pytest_mock import MockerFixture
|
||||
|
||||
import vllm_omni.diffusion.worker.diffusion_worker as worker_module
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _FakeGroup:
|
||||
world_size: int = 1
|
||||
rank_in_group: int = 0
|
||||
|
||||
|
||||
_GROUP_GETTERS = (
|
||||
"get_dp_group",
|
||||
"get_pp_group",
|
||||
"get_sp_group",
|
||||
"get_cfg_group",
|
||||
"get_tp_group",
|
||||
"get_fs_group",
|
||||
"get_hsdp_replicate_group",
|
||||
"get_ep_group",
|
||||
)
|
||||
|
||||
|
||||
def _patch_groups(
|
||||
mocker: MockerFixture,
|
||||
groups: dict[str, _FakeGroup],
|
||||
) -> dict[str, Mock]:
|
||||
patched_getters = {}
|
||||
for getter_name in _GROUP_GETTERS:
|
||||
patched_getters[getter_name] = mocker.patch.object(
|
||||
worker_module,
|
||||
getter_name,
|
||||
return_value=groups.get(getter_name, _FakeGroup()),
|
||||
)
|
||||
return patched_getters
|
||||
|
||||
|
||||
def test_setup_uses_default_name_before_model_parallel_init(
|
||||
mocker: MockerFixture,
|
||||
) -> None:
|
||||
mocker.patch.object(
|
||||
worker_module,
|
||||
"model_parallel_is_initialized",
|
||||
return_value=False,
|
||||
)
|
||||
group_getters = _patch_groups(mocker, {})
|
||||
set_process_title = mocker.patch.object(worker_module, "set_process_title")
|
||||
decorate_logs = mocker.patch.object(worker_module, "decorate_logs")
|
||||
|
||||
worker_module._setup_diffusion_worker_proc_title_and_log_prefix(enable_ep=True, use_hsdp=True)
|
||||
|
||||
set_process_title.assert_called_once_with(
|
||||
name="DiffusionWorker",
|
||||
prefix="vLLM-Omni",
|
||||
)
|
||||
decorate_logs.assert_called_once_with("DiffusionWorker")
|
||||
for group_getter in group_getters.values():
|
||||
group_getter.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("groups", "enable_ep", "use_hsdp", "hsdp_replicate_size", "expected_name"),
|
||||
[
|
||||
({}, False, False, 1, "DiffusionWorker"),
|
||||
({"get_tp_group": _FakeGroup(2, 1)}, False, False, 1, "DiffusionWorker_TP1"),
|
||||
({"get_dp_group": _FakeGroup(4, 2)}, False, False, 1, "DiffusionWorker_DP2"),
|
||||
({"get_pp_group": _FakeGroup(2, 1)}, False, False, 1, "DiffusionWorker_PP1"),
|
||||
(
|
||||
{
|
||||
"get_cfg_group": _FakeGroup(2, 1),
|
||||
"get_tp_group": _FakeGroup(2, 1),
|
||||
},
|
||||
False,
|
||||
False,
|
||||
1,
|
||||
"DiffusionWorker_CFG1_TP1",
|
||||
),
|
||||
(
|
||||
{
|
||||
"get_dp_group": _FakeGroup(2, 1),
|
||||
"get_sp_group": _FakeGroup(4, 2),
|
||||
"get_tp_group": _FakeGroup(2, 1),
|
||||
},
|
||||
False,
|
||||
False,
|
||||
1,
|
||||
"DiffusionWorker_DP1_SP2_TP1",
|
||||
),
|
||||
(
|
||||
{"get_fs_group": _FakeGroup(4, 3)},
|
||||
False,
|
||||
True,
|
||||
1,
|
||||
"DiffusionWorker_FS3",
|
||||
),
|
||||
(
|
||||
{
|
||||
"get_sp_group": _FakeGroup(2, 1),
|
||||
"get_fs_group": _FakeGroup(4, 3),
|
||||
},
|
||||
False,
|
||||
True,
|
||||
1,
|
||||
"DiffusionWorker_SP1_FS3",
|
||||
),
|
||||
(
|
||||
{
|
||||
"get_fs_group": _FakeGroup(1, 0),
|
||||
"get_hsdp_replicate_group": _FakeGroup(2, 1),
|
||||
},
|
||||
False,
|
||||
True,
|
||||
2,
|
||||
"DiffusionWorker_RP1",
|
||||
),
|
||||
({"get_ep_group": _FakeGroup(4, 2)}, True, False, 1, "DiffusionWorker_EP2"),
|
||||
({}, True, False, 1, "DiffusionWorker"),
|
||||
],
|
||||
ids=[
|
||||
"all-singleton",
|
||||
"tp-only",
|
||||
"dp-only",
|
||||
"pp-only",
|
||||
"tp-cfg",
|
||||
"dp-sp-tp",
|
||||
"hsdp-only",
|
||||
"hsdp-sp",
|
||||
"hsdp-replicated",
|
||||
"expert-parallel",
|
||||
"singleton-ep",
|
||||
],
|
||||
)
|
||||
def test_setup_uses_initialized_parallel_groups(
|
||||
mocker: MockerFixture,
|
||||
groups: dict[str, _FakeGroup],
|
||||
enable_ep: bool,
|
||||
use_hsdp: bool,
|
||||
hsdp_replicate_size: int,
|
||||
expected_name: str,
|
||||
) -> None:
|
||||
mocker.patch.object(
|
||||
worker_module,
|
||||
"model_parallel_is_initialized",
|
||||
return_value=True,
|
||||
)
|
||||
group_getters = _patch_groups(mocker, groups)
|
||||
set_process_title = mocker.patch.object(worker_module, "set_process_title")
|
||||
decorate_logs = mocker.patch.object(worker_module, "decorate_logs")
|
||||
|
||||
worker_module._setup_diffusion_worker_proc_title_and_log_prefix(
|
||||
enable_ep=enable_ep,
|
||||
use_hsdp=use_hsdp,
|
||||
hsdp_replicate_size=hsdp_replicate_size,
|
||||
)
|
||||
|
||||
set_process_title.assert_called_once_with(
|
||||
name=expected_name,
|
||||
prefix="vLLM-Omni",
|
||||
)
|
||||
decorate_logs.assert_called_once_with(expected_name)
|
||||
if use_hsdp:
|
||||
group_getters["get_fs_group"].assert_called_once_with()
|
||||
if hsdp_replicate_size > 1:
|
||||
group_getters["get_hsdp_replicate_group"].assert_called_once_with()
|
||||
else:
|
||||
group_getters["get_hsdp_replicate_group"].assert_not_called()
|
||||
else:
|
||||
group_getters["get_fs_group"].assert_not_called()
|
||||
group_getters["get_hsdp_replicate_group"].assert_not_called()
|
||||
if enable_ep:
|
||||
group_getters["get_ep_group"].assert_called_once_with()
|
||||
else:
|
||||
group_getters["get_ep_group"].assert_not_called()
|
||||
|
||||
|
||||
def test_missing_setproctitle_is_non_fatal(mocker: MockerFixture) -> None:
|
||||
mocker.patch.object(
|
||||
worker_module,
|
||||
"model_parallel_is_initialized",
|
||||
return_value=False,
|
||||
)
|
||||
mocker.patch.dict(sys.modules, {"setproctitle": None})
|
||||
decorate_logs = mocker.patch.object(worker_module, "decorate_logs")
|
||||
|
||||
worker_module._setup_diffusion_worker_proc_title_and_log_prefix(enable_ep=False, use_hsdp=False)
|
||||
|
||||
decorate_logs.assert_called_once_with("DiffusionWorker")
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
sys.platform != "linux" or shutil.which("ps") is None or importlib.util.find_spec("setproctitle") is None,
|
||||
reason="requires Linux ps and setproctitle",
|
||||
)
|
||||
def test_process_title_is_visible_through_ps() -> None:
|
||||
expected_name = "vLLM-Omni::DiffusionWorker_TP1"
|
||||
child_script = textwrap.dedent(
|
||||
"""
|
||||
import time
|
||||
|
||||
from vllm.utils.system_utils import set_process_title
|
||||
|
||||
|
||||
set_process_title(
|
||||
name="DiffusionWorker_TP1",
|
||||
prefix="vLLM-Omni",
|
||||
)
|
||||
time.sleep(30)
|
||||
"""
|
||||
)
|
||||
process = subprocess.Popen(
|
||||
[sys.executable, "-c", child_script],
|
||||
stderr=subprocess.PIPE,
|
||||
text=True,
|
||||
)
|
||||
|
||||
observed_title = ""
|
||||
try:
|
||||
for _ in range(100):
|
||||
if process.poll() is not None:
|
||||
stderr = process.stderr.read() if process.stderr is not None else ""
|
||||
pytest.fail(f"child process exited before title check: {stderr}")
|
||||
result = subprocess.run(
|
||||
["ps", "-o", "args=", "-p", str(process.pid)],
|
||||
check=False,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
)
|
||||
observed_title = result.stdout.strip()
|
||||
if expected_name in observed_title:
|
||||
break
|
||||
time.sleep(0.1)
|
||||
else:
|
||||
pytest.fail(f"expected {expected_name!r} in process title, got {observed_title!r}")
|
||||
finally:
|
||||
process.terminate()
|
||||
try:
|
||||
process.wait(timeout=5)
|
||||
except subprocess.TimeoutExpired:
|
||||
process.kill()
|
||||
process.wait(timeout=5)
|
||||
Reference in new issue
Block a user