diff --git a/tests/entrypoints/test_utils.py b/tests/entrypoints/test_utils.py index 4c8522ae0..c6d8e0354 100644 --- a/tests/entrypoints/test_utils.py +++ b/tests/entrypoints/test_utils.py @@ -3,6 +3,8 @@ import os from collections import Counter from dataclasses import dataclass +from pathlib import Path +from types import SimpleNamespace import pytest import torch @@ -13,6 +15,10 @@ 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 +from vllm_omni.engine.stage_init_utils import ( + get_stage_connector_spec, + load_omni_transfer_config_for_model, +) from vllm_omni.entrypoints.utils import ( _convert_dataclasses_to_dict, _filter_dict_like_object, @@ -313,6 +319,58 @@ class TestResolveModelConfigPath: assert result is not None assert "glm_image.yaml" in result + def test_uses_registry_default_deploy_when_hf_model_type_differs( + self, + mocker: MockerFixture, + tmp_path, + ): + from vllm_omni.config.stage_config import _DEPLOY_DIR as _REAL_DEPLOY_DIR + from vllm_omni.entrypoints import utils as utils_mod + + hf_config = SimpleNamespace( + model_type="minicpmo", + architectures=["MiniCPMO"], + version="4.5", + ) + model = "local/minicpmo-registry-default-test" + mocker.patch( + "vllm_omni.entrypoints.utils.get_config", + return_value=hf_config, + ) + mocker.patch( + "vllm_omni.config.config_factory.get_config", + return_value=hf_config, + ) + # A bare deploy/.yaml must not win over the registered + # pipeline default when HF model_type is not an OMNI_PIPELINES key. + (tmp_path / "minicpmo.yaml").write_text("stages: []\n", encoding="utf-8") + (tmp_path / "minicpmo_4_5.yaml").write_bytes((_REAL_DEPLOY_DIR / "minicpmo_4_5.yaml").read_bytes()) + mocker.patch.object(utils_mod, "_DEPLOY_DIR", tmp_path) + + result = resolve_model_config_path(model) + + assert result is not None + assert Path(result).as_posix().endswith("minicpmo_4_5.yaml") + assert Path(result).name == "minicpmo_4_5.yaml" + transfer_config = load_omni_transfer_config_for_model(model, result) + assert transfer_config is not None + + sender = get_stage_connector_spec( + omni_transfer_config=transfer_config, + stage_id=1, + async_chunk=True, + ) + receiver = get_stage_connector_spec( + omni_transfer_config=transfer_config, + stage_id=2, + async_chunk=True, + ) + assert sender["extra"]["role"] == "sender" + assert receiver["extra"]["role"] == "receiver" + assert sender["extra"]["connector_get_sleep_s"] == 0.01 + assert sender["extra"]["connector_get_max_wait_first_chunk"] == 3000 + assert sender["extra"]["connector_get_max_wait"] == 300 + class TestLoadAndResolveStageConfigs: def test_load_and_resolve_with_kwargs(self, mocker: MockerFixture): diff --git a/vllm_omni/entrypoints/utils.py b/vllm_omni/entrypoints/utils.py index 20e9747f9..dab544d62 100644 --- a/vllm_omni/entrypoints/utils.py +++ b/vllm_omni/entrypoints/utils.py @@ -13,6 +13,8 @@ from vllm.transformers_utils.config import get_config, get_hf_file_to_dict from vllm.transformers_utils.repo_utils import file_or_path_exists from vllm_omni.config.config_factory import StageConfigFactory, with_trust_remote_code_override +from vllm_omni.config.pipeline_registry import OMNI_PIPELINES +from vllm_omni.config.stage_config import _DEPLOY_DIR from vllm_omni.config.yaml_util import create_config, load_yaml_config, merge_configs from vllm_omni.diffusion.utils.hf_utils import _looks_like_dreamzero from vllm_omni.entrypoints.stage_utils import _to_dict @@ -226,22 +228,49 @@ def _try_resolve_omni_model_type(model: str) -> str | None: return best_match -def resolve_model_config_path(model: str) -> str: - """Resolve the stage config file path from the model name. +def _registry_default_deploy_path(model: str) -> str | None: + """Return the registered pipeline's default deploy YAML path, if any. - Resolves stage configuration path based on the model type and device type. - First tries to find a device-specific YAML file from stage_configs/{device_type}/ - directory. If not found, falls back to the default config file. + A checkpoint's HF ``model_type`` is not always the registered Omni pipeline + key. MiniCPM-o 4.5, for example, reports ``minicpmo`` while its + predicate-selected pipeline is ``minicpmo_4_5``. Stage construction already + loads that pipeline's default deploy config, so connector discovery must + resolve the same path. + """ + pipeline_config = StageConfigFactory.get_pipeline_config( + model=model, + trust_remote_code=True, + ) + if pipeline_config is None or pipeline_config.default_deploy_config_name is None: + return None + default_deploy_path = _DEPLOY_DIR / pipeline_config.default_deploy_config_name + if default_deploy_path.is_file(): + return str(default_deploy_path) + return None + + +def resolve_model_config_path(model: str) -> str | None: + """Resolve the stage/deploy config file path from the model name. + + Resolves configuration path based on the model type and device type. + Order: + 1. Device-specific stage config under ``stage_configs/{device_type}/`` + 2. If HF ``model_type`` is not an ``OMNI_PIPELINES`` key, the registered + pipeline's ``default_deploy_config_name`` (keeps connectors aligned with + stage construction when HF type and pipeline key differ) + 3. ``deploy/{model_type}.yaml`` + 4. Legacy ``stage_configs/{model_type}.yaml`` + 5. Registered pipeline default deploy config (final fallthrough) Args: model: Model name or path (used to determine model_type) Returns: - String path to the stage configuration file + String path to the stage/deploy configuration file, or ``None`` if no + matching config file exists. Raises: ValueError: If model_type cannot be determined - FileNotFoundError: If no stage config file exists for the model type """ # Try to get config from standard transformers format first try: @@ -296,15 +325,24 @@ def resolve_model_config_path(model: str) -> str: if os.path.exists(complete_config_path): return str(complete_config_path) - deploy_config_path = PROJECT_ROOT / "vllm_omni" / "deploy" / model_type_str + # Prefer the registry default before deploy/.yaml when the + # HF type is not itself a pipeline key. Otherwise a future + # deploy/minicpmo.yaml would desync connectors from stages for MiniCPM-o 4.5. + if normalized_model_type not in OMNI_PIPELINES: + registry_path = _registry_default_deploy_path(model) + if registry_path is not None: + return registry_path + + deploy_config_path = _DEPLOY_DIR / model_type_str if os.path.exists(deploy_config_path): return str(deploy_config_path) stage_config_file = f"vllm_omni/model_executor/stage_configs/{normalized_model_type}.yaml" stage_config_path = PROJECT_ROOT / stage_config_file - if not os.path.exists(stage_config_path): - return None - return str(stage_config_path) + if os.path.exists(stage_config_path): + return str(stage_config_path) + + return _registry_default_deploy_path(model) def load_stage_configs_from_model(