diff --git a/tests/diffusion/test_diffusion_worker_process_title.py b/tests/diffusion/test_diffusion_worker_process_title.py new file mode 100644 index 000000000..092e5fc3d --- /dev/null +++ b/tests/diffusion/test_diffusion_worker_process_title.py @@ -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) diff --git a/vllm_omni/diffusion/worker/diffusion_worker.py b/vllm_omni/diffusion/worker/diffusion_worker.py index f71c2bf8a..1267dd081 100644 --- a/vllm_omni/diffusion/worker/diffusion_worker.py +++ b/vllm_omni/diffusion/worker/diffusion_worker.py @@ -25,10 +25,12 @@ import torch.distributed as dist import zmq from vllm.config import VllmConfig, set_current_vllm_config 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.profiler.wrapper import CudaProfilerWrapper, WorkerProfiler 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.system_utils import decorate_logs, set_process_title from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheSpec from vllm.v1.worker.utils import request_memory from vllm.v1.worker.workspace import init_workspace_manager @@ -46,8 +48,14 @@ from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode from vllm_omni.diffusion.diffusion_kv.metadata import DiffusionKVMetadata from vllm_omni.diffusion.distributed.parallel_state import ( destroy_distributed_env, + get_cfg_group, + get_dp_group, + get_fs_group, + get_pp_group, + get_sp_group, init_distributed_environment, initialize_model_parallel, + model_parallel_is_initialized, ) from vllm_omni.diffusion.forward_context import set_forward_context from vllm_omni.diffusion.ipc import DIFFUSION_RPC_RESULT_ENVELOPE, pack_diffusion_output_shm @@ -102,6 +110,38 @@ def _all_gather_rank_values(value: Any) -> list[Any]: 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]: """Run one rank-local probe without stranding peers on local failure.""" @@ -297,6 +337,7 @@ class DiffusionWorker: enable_expert_parallel=parallel_config.enable_expert_parallel, use_hsdp=parallel_config.use_hsdp, ) + _setup_diffusion_worker_proc_title_and_log_prefix(enable_ep=parallel_config.enable_expert_parallel) if ( getattr(self.od_config, "diffusion_kv_mode", DiffusionKVCacheMode.DENSE_LEGACY) is DiffusionKVCacheMode.PAGED_SCHEDULER @@ -1399,13 +1440,7 @@ class WorkerProc: set_death_signal(signal.SIGTERM) - # 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 + _setup_diffusion_worker_proc_title_and_log_prefix(enable_ep=od_config.parallel_config.enable_expert_parallel) load_omni_general_plugins() worker_proc = None