[Core] Make diffusion worker titles topology-aware (#6307)
Signed-off-by: Dragonborn <1078757180@qq.com>
This commit is contained in:
prettygirlisnotme authored and GitHub committed
1 parent
009b992c5b
commit
13173d6bbb
2 files changed
+260
-7
No files matched your search
@@ -0,0 +1,218 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM 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_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)
|
||||
|
||||
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", "expected_name"),
|
||||
[
|
||||
({}, False, "DiffusionWorker"),
|
||||
({"get_tp_group": _FakeGroup(2, 1)}, False, "DiffusionWorker_TP1"),
|
||||
({"get_dp_group": _FakeGroup(4, 2)}, False, "DiffusionWorker_DP2"),
|
||||
({"get_pp_group": _FakeGroup(2, 1)}, False, "DiffusionWorker_PP1"),
|
||||
(
|
||||
{
|
||||
"get_cfg_group": _FakeGroup(2, 1),
|
||||
"get_tp_group": _FakeGroup(2, 1),
|
||||
},
|
||||
False,
|
||||
"DiffusionWorker_CFG1_TP1",
|
||||
),
|
||||
(
|
||||
{
|
||||
"get_dp_group": _FakeGroup(2, 1),
|
||||
"get_sp_group": _FakeGroup(4, 2),
|
||||
"get_tp_group": _FakeGroup(2, 1),
|
||||
},
|
||||
False,
|
||||
"DiffusionWorker_DP1_SP2_TP1",
|
||||
),
|
||||
(
|
||||
{
|
||||
"get_dp_group": _FakeGroup(4, 3),
|
||||
"get_fs_group": _FakeGroup(2, 1),
|
||||
},
|
||||
False,
|
||||
"DiffusionWorker_DP3_FS1",
|
||||
),
|
||||
(
|
||||
{"get_ep_group": _FakeGroup(4, 2)},
|
||||
True,
|
||||
"DiffusionWorker_EP2",
|
||||
),
|
||||
({}, True, "DiffusionWorker"),
|
||||
],
|
||||
ids=[
|
||||
"all-singleton",
|
||||
"tp-only",
|
||||
"dp-only",
|
||||
"pp-only",
|
||||
"tp-cfg",
|
||||
"dp-sp-tp",
|
||||
"standalone-hsdp",
|
||||
"expert-parallel",
|
||||
"singleton-ep",
|
||||
],
|
||||
)
|
||||
def test_setup_uses_initialized_parallel_groups(
|
||||
mocker: MockerFixture,
|
||||
groups: dict[str, _FakeGroup],
|
||||
enable_ep: bool,
|
||||
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)
|
||||
|
||||
set_process_title.assert_called_once_with(
|
||||
name=expected_name,
|
||||
prefix="vLLM-Omni",
|
||||
)
|
||||
decorate_logs.assert_called_once_with(expected_name)
|
||||
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)
|
||||
|
||||
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