From d5aced92966e6f12134fe9a30679fdeab07e4bb9 Mon Sep 17 00:00:00 2001 From: Hongsheng Liu Date: Fri, 28 Aug 2026 17:30:03 +0800 Subject: [PATCH] Revert "[Core] Make diffusion worker titles topology-aware" (#6717) --- .../test_diffusion_worker_process_title.py | 218 ------------------ .../diffusion/worker/diffusion_worker.py | 49 +--- 2 files changed, 7 insertions(+), 260 deletions(-) delete mode 100644 tests/diffusion/test_diffusion_worker_process_title.py diff --git a/tests/diffusion/test_diffusion_worker_process_title.py b/tests/diffusion/test_diffusion_worker_process_title.py deleted file mode 100644 index 092e5fc3d..000000000 --- a/tests/diffusion/test_diffusion_worker_process_title.py +++ /dev/null @@ -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) diff --git a/vllm_omni/diffusion/worker/diffusion_worker.py b/vllm_omni/diffusion/worker/diffusion_worker.py index 1267dd081..f71c2bf8a 100644 --- a/vllm_omni/diffusion/worker/diffusion_worker.py +++ b/vllm_omni/diffusion/worker/diffusion_worker.py @@ -25,12 +25,10 @@ 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 @@ -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.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 @@ -110,38 +102,6 @@ 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.""" @@ -337,7 +297,6 @@ 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 @@ -1440,7 +1399,13 @@ class WorkerProc: 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() worker_proc = None