[Entrypoint][Refactor]Stage CLI Refactor (#2020)
Signed-off-by: wuhang <wuhang6@huawei.com> Signed-off-by: Didan Deng <33117903+wtomin@users.noreply.github.com> Co-authored-by: Didan Deng <33117903+wtomin@users.noreply.github.com>
This commit is contained in:
wuhangandDidan Deng authored and GitHub committed
1 parent
f3f2dc590c
commit
cb91cbe61e
19 files changed
+3773
-273
No files matched your search
@@ -0,0 +1,195 @@
|
||||
"""Unit tests for the Omni serve CLI helpers."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
from unittest.mock import Mock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm_omni.entrypoints.cli.serve import run_headless
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
def _make_headless_args() -> argparse.Namespace:
|
||||
return argparse.Namespace(
|
||||
model="fake-model",
|
||||
stage_id=3,
|
||||
omni_master_address="127.0.0.1",
|
||||
omni_master_port=26000,
|
||||
api_server_count=0,
|
||||
worker_backend="multi_process",
|
||||
stage_configs_path=None,
|
||||
log_stats=False,
|
||||
disable_log_stats=False,
|
||||
)
|
||||
|
||||
|
||||
def test_run_headless_registers_stage_once_and_launches_all_local_engines() -> None:
|
||||
args = _make_headless_args()
|
||||
stage_cfg = Mock(stage_id=3)
|
||||
stage_cfgs = [stage_cfg]
|
||||
parallel_config = Mock(
|
||||
data_parallel_size_local=2,
|
||||
data_parallel_rank=4,
|
||||
data_parallel_rank_local=1,
|
||||
node_rank_within_dp=0,
|
||||
)
|
||||
vllm_config = Mock(parallel_config=parallel_config)
|
||||
executor_class = Mock()
|
||||
engine_manager = Mock()
|
||||
|
||||
with (
|
||||
patch(
|
||||
"vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
|
||||
return_value=("/fake/stages.yaml", stage_cfgs),
|
||||
),
|
||||
patch("vllm_omni.engine.stage_init_utils.prepare_engine_environment"),
|
||||
patch("vllm_omni.engine.stage_init_utils.load_omni_transfer_config_for_model", return_value=Mock()),
|
||||
patch("vllm_omni.engine.stage_init_utils.get_stage_connector_spec", return_value={}),
|
||||
patch("vllm_omni.engine.stage_init_utils.build_engine_args_dict", return_value={}),
|
||||
patch(
|
||||
"vllm_omni.distributed.omni_connectors.utils.initialization.resolve_omni_kv_config_for_stage",
|
||||
return_value=(None, None, None),
|
||||
),
|
||||
patch(
|
||||
"vllm_omni.engine.stage_init_utils.build_vllm_config",
|
||||
return_value=(vllm_config, executor_class),
|
||||
) as mock_build_vllm_config,
|
||||
patch(
|
||||
"vllm_omni.engine.stage_engine_startup.register_stage_with_omni_master",
|
||||
return_value="tcp://127.0.0.1:26001",
|
||||
) as mock_register,
|
||||
patch("vllm.v1.engine.utils.CoreEngineProcManager", return_value=engine_manager) as mock_manager_cls,
|
||||
patch("signal.signal"),
|
||||
):
|
||||
run_headless(args)
|
||||
|
||||
mock_build_vllm_config.assert_called_once_with(
|
||||
stage_cfg,
|
||||
"fake-model",
|
||||
stage_connector_spec={},
|
||||
engine_args_dict={},
|
||||
headless=True,
|
||||
)
|
||||
mock_register.assert_called_once_with(
|
||||
omni_master_address="127.0.0.1",
|
||||
omni_master_port=26000,
|
||||
omni_stage_id=3,
|
||||
omni_stage_config=stage_cfg,
|
||||
coordinator=None,
|
||||
)
|
||||
mock_manager_cls.assert_called_once()
|
||||
manager_kwargs = mock_manager_cls.call_args.kwargs
|
||||
assert manager_kwargs["local_engine_count"] == 2
|
||||
assert manager_kwargs["start_index"] == 4
|
||||
assert manager_kwargs["local_start_index"] == 0
|
||||
assert manager_kwargs["local_client"] is False
|
||||
assert manager_kwargs["handshake_address"] == "tcp://127.0.0.1:26001"
|
||||
assert manager_kwargs["log_stats"] is False
|
||||
engine_manager.join_first.assert_called_once_with()
|
||||
engine_manager.shutdown.assert_called_once_with()
|
||||
|
||||
|
||||
def test_run_headless_honors_explicit_log_stats_flag() -> None:
|
||||
args = _make_headless_args()
|
||||
args.log_stats = True
|
||||
stage_cfg = Mock(stage_id=3)
|
||||
stage_cfgs = [stage_cfg]
|
||||
parallel_config = Mock(
|
||||
data_parallel_size_local=2,
|
||||
data_parallel_rank=4,
|
||||
data_parallel_rank_local=1,
|
||||
node_rank_within_dp=0,
|
||||
)
|
||||
vllm_config = Mock(parallel_config=parallel_config)
|
||||
executor_class = Mock()
|
||||
engine_manager = Mock()
|
||||
|
||||
with (
|
||||
patch(
|
||||
"vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
|
||||
return_value=("/fake/stages.yaml", stage_cfgs),
|
||||
),
|
||||
patch("vllm_omni.engine.stage_init_utils.prepare_engine_environment"),
|
||||
patch("vllm_omni.engine.stage_init_utils.load_omni_transfer_config_for_model", return_value=Mock()),
|
||||
patch("vllm_omni.engine.stage_init_utils.get_stage_connector_spec", return_value={}),
|
||||
patch("vllm_omni.engine.stage_init_utils.build_engine_args_dict", return_value={}),
|
||||
patch(
|
||||
"vllm_omni.distributed.omni_connectors.utils.initialization.resolve_omni_kv_config_for_stage",
|
||||
return_value=(None, None, None),
|
||||
),
|
||||
patch(
|
||||
"vllm_omni.engine.stage_init_utils.build_vllm_config",
|
||||
return_value=(vllm_config, executor_class),
|
||||
),
|
||||
patch(
|
||||
"vllm_omni.engine.stage_engine_startup.register_stage_with_omni_master",
|
||||
return_value="tcp://127.0.0.1:26001",
|
||||
),
|
||||
patch("vllm.v1.engine.utils.CoreEngineProcManager", return_value=engine_manager) as mock_manager_cls,
|
||||
patch("signal.signal"),
|
||||
):
|
||||
run_headless(args)
|
||||
|
||||
manager_kwargs = mock_manager_cls.call_args.kwargs
|
||||
assert manager_kwargs["log_stats"] is True
|
||||
|
||||
|
||||
def test_run_headless_launches_diffusion_stage_via_omni_master() -> None:
|
||||
args = _make_headless_args()
|
||||
stage_cfg = Mock(stage_id=3, stage_type="diffusion")
|
||||
stage_cfg.engine_args = Mock()
|
||||
stage_cfg.engine_input_source = []
|
||||
stage_cfgs = [stage_cfg]
|
||||
metadata = Mock(stage_id=3)
|
||||
od_config = Mock()
|
||||
proc = Mock()
|
||||
proc.exitcode = 0
|
||||
proc.is_alive.return_value = False
|
||||
|
||||
with (
|
||||
patch(
|
||||
"vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
|
||||
return_value=("/fake/stages.yaml", stage_cfgs),
|
||||
),
|
||||
patch("vllm_omni.engine.stage_init_utils.prepare_engine_environment"),
|
||||
patch("vllm_omni.engine.stage_init_utils.load_omni_transfer_config_for_model", return_value=Mock()),
|
||||
patch(
|
||||
"vllm_omni.distributed.omni_connectors.utils.initialization.resolve_omni_kv_config_for_stage",
|
||||
return_value=(None, None, None),
|
||||
),
|
||||
patch("vllm_omni.engine.stage_init_utils.extract_stage_metadata", return_value=metadata),
|
||||
patch("vllm_omni.engine.stage_init_utils.inject_kv_stage_info") as mock_inject_stage_info,
|
||||
patch("vllm_omni.engine.stage_init_utils.build_diffusion_config", return_value=od_config),
|
||||
patch(
|
||||
"vllm_omni.engine.stage_engine_startup.register_stage_with_omni_master",
|
||||
return_value=("tcp://127.0.0.1:26001", "tcp://127.0.0.1:26002", "tcp://127.0.0.1:26003"),
|
||||
) as mock_register,
|
||||
patch(
|
||||
"vllm_omni.diffusion.stage_diffusion_proc.spawn_diffusion_proc",
|
||||
return_value=(proc, "tcp://127.0.0.1:26001", "tcp://127.0.0.1:26002", "tcp://127.0.0.1:26003"),
|
||||
) as mock_spawn,
|
||||
patch("vllm_omni.diffusion.stage_diffusion_proc.complete_diffusion_handshake") as mock_handshake,
|
||||
patch("signal.signal"),
|
||||
):
|
||||
run_headless(args)
|
||||
|
||||
mock_inject_stage_info.assert_called_once_with(stage_cfg, 3)
|
||||
mock_register.assert_called_once_with(
|
||||
omni_master_address="127.0.0.1",
|
||||
omni_master_port=26000,
|
||||
omni_stage_id=3,
|
||||
omni_stage_config=stage_cfg,
|
||||
return_addresses=True,
|
||||
)
|
||||
mock_spawn.assert_called_once_with(
|
||||
"fake-model",
|
||||
od_config,
|
||||
handshake_address="tcp://127.0.0.1:26001",
|
||||
request_address="tcp://127.0.0.1:26002",
|
||||
response_address="tcp://127.0.0.1:26003",
|
||||
)
|
||||
mock_handshake.assert_called_once_with(proc, "tcp://127.0.0.1:26001")
|
||||
proc.join.assert_called_once_with()
|
||||
@@ -8,6 +8,7 @@ import pytest
|
||||
import torch
|
||||
from pytest_mock import MockerFixture
|
||||
|
||||
from vllm_omni.config.yaml_util import create_config
|
||||
from vllm_omni.diffusion.data import OmniDiffusionConfig
|
||||
from vllm_omni.engine.arg_utils import OmniEngineArgs
|
||||
from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
|
||||
@@ -16,6 +17,7 @@ from vllm_omni.entrypoints.utils import (
|
||||
_filter_dict_like_object,
|
||||
filter_dataclass_kwargs,
|
||||
load_and_resolve_stage_configs,
|
||||
load_stage_configs_from_yaml,
|
||||
resolve_model_config_path,
|
||||
)
|
||||
|
||||
@@ -322,3 +324,70 @@ class TestLoadAndResolveStageConfigs:
|
||||
assert config_path is None
|
||||
assert len(stage_configs) == 1
|
||||
assert "dtype" in stage_configs[0]["engine_args"]
|
||||
|
||||
|
||||
class TestLoadStageConfigsFromYaml:
|
||||
"""Regression tests for stage-config loading and merging."""
|
||||
|
||||
def test_deep_merges_stage_engine_args(self, mocker: MockerFixture):
|
||||
yaml_config = create_config(
|
||||
{
|
||||
"async_chunk": True,
|
||||
"stage_args": [
|
||||
{
|
||||
"stage_id": 0,
|
||||
"runtime": {"device": 0},
|
||||
"engine_args": {
|
||||
"parallel_config": {"tensor_parallel_size": 4},
|
||||
},
|
||||
}
|
||||
],
|
||||
}
|
||||
)
|
||||
mocker.patch(
|
||||
"vllm_omni.entrypoints.utils.load_yaml_config",
|
||||
return_value=yaml_config,
|
||||
)
|
||||
|
||||
stages = load_stage_configs_from_yaml(
|
||||
"fake.yaml",
|
||||
base_engine_args={
|
||||
"parallel_config": {
|
||||
"tensor_parallel_size": 1,
|
||||
"pipeline_parallel_size": 2,
|
||||
},
|
||||
"model": "base-model",
|
||||
},
|
||||
)
|
||||
|
||||
merged_engine_args = stages[0]["engine_args"]
|
||||
assert merged_engine_args["parallel_config"]["tensor_parallel_size"] == 4
|
||||
assert merged_engine_args["parallel_config"]["pipeline_parallel_size"] == 2
|
||||
assert merged_engine_args["model"] == "base-model"
|
||||
assert merged_engine_args["async_chunk"] is True
|
||||
|
||||
def test_merges_nested_stage_engine_args(self, mocker: MockerFixture):
|
||||
yaml_config = create_config(
|
||||
{
|
||||
"stage_args": [
|
||||
{
|
||||
"stage_id": 0,
|
||||
"engine_args": {
|
||||
"nested": {"override": 2},
|
||||
},
|
||||
}
|
||||
],
|
||||
}
|
||||
)
|
||||
mocker.patch(
|
||||
"vllm_omni.entrypoints.utils.load_yaml_config",
|
||||
return_value=yaml_config,
|
||||
)
|
||||
|
||||
stages = load_stage_configs_from_yaml(
|
||||
"fake.yaml",
|
||||
base_engine_args={"nested": {"base": 1}},
|
||||
)
|
||||
|
||||
assert stages[0]["engine_args"]["nested"]["base"] == 1
|
||||
assert stages[0]["engine_args"]["nested"]["override"] == 2
|
||||
Reference in new issue
Block a user