vllm-omni/tests/diffusion/test_diffusion_worker_process_title.py

# 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)