Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
dc31c49b76 |
No files matched your search
@@ -1,218 +0,0 @@
|
|||||||
# 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)
|
|
||||||
@@ -25,12 +25,10 @@ import torch.distributed as dist
|
|||||||
import zmq
|
import zmq
|
||||||
from vllm.config import VllmConfig, set_current_vllm_config
|
from vllm.config import VllmConfig, set_current_vllm_config
|
||||||
from vllm.distributed.device_communicators.shm_broadcast import MessageQueue
|
from vllm.distributed.device_communicators.shm_broadcast import MessageQueue
|
||||||
from vllm.distributed.parallel_state import get_ep_group, get_tp_group
|
|
||||||
from vllm.logger import init_logger
|
from vllm.logger import init_logger
|
||||||
from vllm.profiler.wrapper import CudaProfilerWrapper, WorkerProfiler
|
from vllm.profiler.wrapper import CudaProfilerWrapper, WorkerProfiler
|
||||||
from vllm.utils.import_utils import resolve_obj_by_qualname
|
from vllm.utils.import_utils import resolve_obj_by_qualname
|
||||||
from vllm.utils.mem_utils import GiB_bytes, MemorySnapshot, format_gib, memory_profiling
|
from vllm.utils.mem_utils import GiB_bytes, MemorySnapshot, format_gib, memory_profiling
|
||||||
from vllm.utils.system_utils import decorate_logs, set_process_title
|
|
||||||
from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheSpec
|
from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheSpec
|
||||||
from vllm.v1.worker.utils import request_memory
|
from vllm.v1.worker.utils import request_memory
|
||||||
from vllm.v1.worker.workspace import init_workspace_manager
|
from vllm.v1.worker.workspace import init_workspace_manager
|
||||||
@@ -48,14 +46,8 @@ from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
|
|||||||
from vllm_omni.diffusion.diffusion_kv.metadata import DiffusionKVMetadata
|
from vllm_omni.diffusion.diffusion_kv.metadata import DiffusionKVMetadata
|
||||||
from vllm_omni.diffusion.distributed.parallel_state import (
|
from vllm_omni.diffusion.distributed.parallel_state import (
|
||||||
destroy_distributed_env,
|
destroy_distributed_env,
|
||||||
get_cfg_group,
|
|
||||||
get_dp_group,
|
|
||||||
get_fs_group,
|
|
||||||
get_pp_group,
|
|
||||||
get_sp_group,
|
|
||||||
init_distributed_environment,
|
init_distributed_environment,
|
||||||
initialize_model_parallel,
|
initialize_model_parallel,
|
||||||
model_parallel_is_initialized,
|
|
||||||
)
|
)
|
||||||
from vllm_omni.diffusion.forward_context import set_forward_context
|
from vllm_omni.diffusion.forward_context import set_forward_context
|
||||||
from vllm_omni.diffusion.ipc import DIFFUSION_RPC_RESULT_ENVELOPE, pack_diffusion_output_shm
|
from vllm_omni.diffusion.ipc import DIFFUSION_RPC_RESULT_ENVELOPE, pack_diffusion_output_shm
|
||||||
@@ -110,38 +102,6 @@ def _all_gather_rank_values(value: Any) -> list[Any]:
|
|||||||
return values
|
return values
|
||||||
|
|
||||||
|
|
||||||
def _setup_diffusion_worker_proc_title_and_log_prefix(enable_ep: bool) -> None:
|
|
||||||
"""Set the worker process title and log prefix from initialized groups."""
|
|
||||||
process_name = "DiffusionWorker"
|
|
||||||
if model_parallel_is_initialized():
|
|
||||||
dp_group = get_dp_group()
|
|
||||||
pp_group = get_pp_group()
|
|
||||||
sp_group = get_sp_group()
|
|
||||||
cfg_group = get_cfg_group()
|
|
||||||
tp_group = get_tp_group()
|
|
||||||
fs_group = get_fs_group()
|
|
||||||
|
|
||||||
if dp_group.world_size > 1:
|
|
||||||
process_name += f"_DP{dp_group.rank_in_group}"
|
|
||||||
if pp_group.world_size > 1:
|
|
||||||
process_name += f"_PP{pp_group.rank_in_group}"
|
|
||||||
if sp_group.world_size > 1:
|
|
||||||
process_name += f"_SP{sp_group.rank_in_group}"
|
|
||||||
if cfg_group.world_size > 1:
|
|
||||||
process_name += f"_CFG{cfg_group.rank_in_group}"
|
|
||||||
if tp_group.world_size > 1:
|
|
||||||
process_name += f"_TP{tp_group.rank_in_group}"
|
|
||||||
if fs_group.world_size > 1:
|
|
||||||
process_name += f"_FS{fs_group.rank_in_group}"
|
|
||||||
if enable_ep:
|
|
||||||
ep_group = get_ep_group()
|
|
||||||
if ep_group.world_size > 1:
|
|
||||||
process_name += f"_EP{ep_group.rank_in_group}"
|
|
||||||
|
|
||||||
set_process_title(name=process_name, prefix="vLLM-Omni")
|
|
||||||
decorate_logs(process_name)
|
|
||||||
|
|
||||||
|
|
||||||
def _run_and_gather_rank_values(operation: str, func: Callable[[], Any]) -> list[Any]:
|
def _run_and_gather_rank_values(operation: str, func: Callable[[], Any]) -> list[Any]:
|
||||||
"""Run one rank-local probe without stranding peers on local failure."""
|
"""Run one rank-local probe without stranding peers on local failure."""
|
||||||
|
|
||||||
@@ -337,7 +297,6 @@ class DiffusionWorker:
|
|||||||
enable_expert_parallel=parallel_config.enable_expert_parallel,
|
enable_expert_parallel=parallel_config.enable_expert_parallel,
|
||||||
use_hsdp=parallel_config.use_hsdp,
|
use_hsdp=parallel_config.use_hsdp,
|
||||||
)
|
)
|
||||||
_setup_diffusion_worker_proc_title_and_log_prefix(enable_ep=parallel_config.enable_expert_parallel)
|
|
||||||
if (
|
if (
|
||||||
getattr(self.od_config, "diffusion_kv_mode", DiffusionKVCacheMode.DENSE_LEGACY)
|
getattr(self.od_config, "diffusion_kv_mode", DiffusionKVCacheMode.DENSE_LEGACY)
|
||||||
is DiffusionKVCacheMode.PAGED_SCHEDULER
|
is DiffusionKVCacheMode.PAGED_SCHEDULER
|
||||||
@@ -1440,7 +1399,13 @@ class WorkerProc:
|
|||||||
|
|
||||||
set_death_signal(signal.SIGTERM)
|
set_death_signal(signal.SIGTERM)
|
||||||
|
|
||||||
_setup_diffusion_worker_proc_title_and_log_prefix(enable_ep=od_config.parallel_config.enable_expert_parallel)
|
# Set process title for visibility in nvidia-smi / htop (optional, non-fatal)
|
||||||
|
try:
|
||||||
|
import setproctitle
|
||||||
|
|
||||||
|
setproctitle.setproctitle(f"vLLM-Omni::DiffusionWorker-{rank}")
|
||||||
|
except ImportError:
|
||||||
|
pass # setproctitle not installed, skip process title setting
|
||||||
|
|
||||||
load_omni_general_plugins()
|
load_omni_general_plugins()
|
||||||
worker_proc = None
|
worker_proc = None
|
||||||
|
|||||||
Reference in new issue
Block a user