From 6f36fd3b7838c4de3a4b9d4de90c3bba7776c91e Mon Sep 17 00:00:00 2001 From: Ma1oneZhang <73012001+Ma1oneZhang@users.noreply.github.com> Date: Sat, 15 Aug 2026 14:11:48 +0800 Subject: [PATCH] [Bugfix] Preserve object storage URIs during model resolution (#5036) Signed-off-by: Ziyang Zhang --- tests/config/test_config_factory.py | 90 ++++++++++++++++++- .../test_omni_snapshot_download.py | 60 +++++++++++++ tests/entrypoints/test_utils.py | 43 +++++++++ vllm_omni/config/config_factory.py | 48 ++++++++-- vllm_omni/entrypoints/omni_base.py | 9 ++ vllm_omni/entrypoints/utils.py | 35 +++++--- 6 files changed, 267 insertions(+), 18 deletions(-) diff --git a/tests/config/test_config_factory.py b/tests/config/test_config_factory.py index 9708c1c50..5bfbb67e8 100644 --- a/tests/config/test_config_factory.py +++ b/tests/config/test_config_factory.py @@ -14,7 +14,8 @@ import pytest from transformers import PretrainedConfig, Qwen3OmniMoeConfig from tests.helpers.stage_config import get_deploy_config_path, get_deploy_config_stage -from vllm_omni.config.config_factory import StageConfigFactory +from vllm_omni.config import config_factory as config_factory_module +from vllm_omni.config.config_factory import StageConfigFactory, _materialize_object_storage_configs from vllm_omni.config.endpoint_policy import EndpointRestriction, OmniServingCapability from vllm_omni.config.omni_config import VllmOmniConfig from vllm_omni.config.pipeline_registry import OMNI_PIPELINES, register_pipeline, resolve_pipeline_config @@ -55,6 +56,7 @@ def clear_config_factory_caches(): yield StageConfigFactory.get_hf_config.cache_clear() StageConfigFactory.try_infer_model_type.cache_clear() + _materialize_object_storage_configs.cache_clear() Q3_OMNI_ALL_STAGES_HF_CONFIG = Qwen3OmniMoeConfig(enable_audio_output=True) @@ -2765,3 +2767,89 @@ class TestPipelineConfigResolvers: pass assert resolver(NotTheRightHfConfig()) is None + + +class TestObjectStorageConfigResolution: + """Object-storage URIs must be materialized locally before any HF-style reads. + + Regression coverage for review on the Run:AI PR: ``get_config`` and the + config.json/model_index.json fallbacks crashed with ``HFValidationError`` + for ``s3://`` URIs, and the name-match fallback scanned the whole URI, so + a bucket named after another pipeline could hijack pipeline selection. + """ + + @pytest.fixture + def fake_object_storage(self, monkeypatch, tmp_path): + """Replace ObjectStorageModel with a fake that "pulls" into a tmp dir. + + Returns a recorder: calls to the returned callable write files into + the materialized directory, and ``recorder.pulls`` records every + ``pull_files`` invocation. + """ + pulls: list[tuple[str, list[str] | None, list[str] | None]] = [] + materialized_files: list[tuple[str, str]] = [] + + class FakeObjectStorageModel: + def __init__(self, url: str): + self.dir = str(tmp_path) + + def pull_files(self, model_path, allow_pattern=None, ignore_pattern=None): + pulls.append((model_path, allow_pattern, ignore_pattern)) + for name, content in materialized_files: + (tmp_path / name).write_text(content) + + monkeypatch.setattr(config_factory_module, "ObjectStorageModel", FakeObjectStorageModel) + + class Recorder: + def set_files(self, files: list[tuple[str, str]]) -> None: + materialized_files.clear() + materialized_files.extend(files) + + recorder = Recorder() + recorder.pulls = pulls + return recorder + + def test_passthrough_for_non_uri(self, fake_object_storage): + assert _materialize_object_storage_configs("org/model") == "org/model" + assert _materialize_object_storage_configs("/local/model") == "/local/model" + assert fake_object_storage.pulls == [] + + def test_materialize_pulls_configs_once_per_uri(self, fake_object_storage): + uri = "s3://bucket/model" + + first = _materialize_object_storage_configs(uri) + second = _materialize_object_storage_configs(uri) + + assert first == second + assert len(fake_object_storage.pulls) == 1 + pulled_uri, allow_pattern, ignore_pattern = fake_object_storage.pulls[0] + assert pulled_uri == uri + assert "*.json" in allow_pattern + assert ignore_pattern is None + + def test_try_infer_model_type_reads_materialized_config(self, fake_object_storage): + """A misleading bucket name must not beat the real config.json content.""" + fake_object_storage.set_files( + [ + ( + "config.json", + '{"model_type": "qwen3_omni_moe", "architectures": ["Qwen3OmniMoeForConditionalGeneration"]}', + ) + ] + ) + uri = "s3://qwen3-tts-models/Qwen3-Omni-30B-A3B-Instruct" + + model = StageConfigFactory.try_infer_model_type(model=uri, trust_remote_code=False) + + assert model == "qwen3_omni_moe" + assert fake_object_storage.pulls # materialization happened + + def test_name_match_fallback_scans_basename_only(self, fake_object_storage): + # Empty config.json (CosyVoice3 style) forces the name-match fallback. + fake_object_storage.set_files([("config.json", "{}")]) + + deceptive = "s3://qwen3-tts-models/plain-checkpoint" + assert StageConfigFactory.try_infer_model_type(model=deceptive, trust_remote_code=False) is None + + matching = "s3://any-bucket/my-cosyvoice3-model" + assert StageConfigFactory.try_infer_model_type(model=matching, trust_remote_code=False) == "cosyvoice3" diff --git a/tests/entrypoints/test_omni_snapshot_download.py b/tests/entrypoints/test_omni_snapshot_download.py index d2c76a249..58dd1d19d 100644 --- a/tests/entrypoints/test_omni_snapshot_download.py +++ b/tests/entrypoints/test_omni_snapshot_download.py @@ -10,8 +10,10 @@ ModelScope path. import sys import types +from pathlib import Path import pytest +from pytest_mock import MockerFixture from vllm import envs from vllm_omni.entrypoints import omni_base @@ -21,6 +23,64 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu] MODEL_ID = "Qwen/Qwen3-Omni-30B-A3B-Instruct" +@pytest.mark.parametrize( + "model_uri", + [ + "s3://bucket/model", + "gs://bucket/model", + "az://bucket/model", + ], +) +def test_omni_snapshot_download_preserves_object_storage_uri( + model_uri: str, + mocker: MockerFixture, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.delenv("VLLM_USE_MODELSCOPE", raising=False) + hf_download = mocker.patch.object(omni_base, "download_weights_from_hf_specific") + + assert omni_base.omni_snapshot_download(model_uri) == model_uri + hf_download.assert_not_called() + + +def test_omni_snapshot_download_preserves_existing_local_path( + tmp_path: Path, + mocker: MockerFixture, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.delenv("VLLM_USE_MODELSCOPE", raising=False) + absolute_model_path = tmp_path / "absolute-model" + absolute_model_path.mkdir() + relative_model_path = "relative-model" + (tmp_path / relative_model_path).mkdir() + monkeypatch.chdir(tmp_path) + hf_download = mocker.patch.object(omni_base, "download_weights_from_hf_specific") + + assert omni_base.omni_snapshot_download(str(absolute_model_path)) == str(absolute_model_path) + assert omni_base.omni_snapshot_download(relative_model_path) == relative_model_path + hf_download.assert_not_called() + + +def test_omni_snapshot_download_uses_hf_for_relative_repo_id( + mocker: MockerFixture, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.delenv("VLLM_USE_MODELSCOPE", raising=False) + model_id = "org/model" + # Stub the modular-index probe so the test never reaches huggingface.co. + hf_file_probe = mocker.patch.object(omni_base, "file_or_path_exists", return_value=False) + hf_download = mocker.patch.object(omni_base, "download_weights_from_hf_specific") + + assert omni_base.omni_snapshot_download(model_id) == model_id + hf_file_probe.assert_called_once_with(model_id, "modular_model_index.json", revision=None) + hf_download.assert_called_once_with( + model_name_or_path=model_id, + cache_dir=None, + allow_patterns=["*"], + require_all=True, + ) + + @pytest.fixture def download_backend(monkeypatch: pytest.MonkeyPatch): """Run ``omni_snapshot_download`` and report which backend it selected. diff --git a/tests/entrypoints/test_utils.py b/tests/entrypoints/test_utils.py index a20c97a58..ccc6f69e8 100644 --- a/tests/entrypoints/test_utils.py +++ b/tests/entrypoints/test_utils.py @@ -22,6 +22,7 @@ from vllm_omni.engine.stage_init_utils import ( from vllm_omni.entrypoints.utils import ( _convert_dataclasses_to_dict, _filter_dict_like_object, + _try_resolve_omni_model_type, coerce_param_message_types, filter_dataclass_kwargs, filter_stages, @@ -371,6 +372,48 @@ class TestResolveModelConfigPath: assert sender["extra"]["connector_get_max_wait_first_chunk"] == 3000 assert sender["extra"]["connector_get_max_wait"] == 300 + def test_object_storage_uri_resolves_via_materialized_configs( + self, + mocker: MockerFixture, + tmp_path, + ): + """An object-storage URI must not reach HF helpers; resolution reads the + locally materialized config files instead.""" + hf_config = SimpleNamespace( + model_type="qwen3_omni_moe", + architectures=["Qwen3OmniMoeForConditionalGeneration"], + ) + uri = "s3://qwen3-tts-models/Qwen3-Omni-30B-A3B-Instruct" + materialize = mocker.patch( + "vllm_omni.entrypoints.utils._materialize_object_storage_configs", + return_value=str(tmp_path), + ) + get_config = mocker.patch( + "vllm_omni.entrypoints.utils.get_config", + return_value=hf_config, + ) + + result = resolve_model_config_path(uri) + + materialize.assert_called_once_with(uri) + get_config.assert_called_once() + assert get_config.call_args[0][0] == str(tmp_path) + assert result is not None + assert Path(result).as_posix().endswith("qwen3_omni_moe.yaml") + + +class TestTryResolveOmniModelType: + """Name matching must only scan the last path component.""" + + def test_uri_bucket_name_cannot_hijack_pipeline(self): + # "qwen3-tts-models" normalizes to "qwen3ttsmodels", which contains the + # registered key "qwen3tts" — only the basename may participate, so a + # misleadingly named bucket selects nothing. + assert _try_resolve_omni_model_type("s3://qwen3-tts-models/plain-checkpoint") is None + + def test_basename_still_matches_pipeline(self): + assert _try_resolve_omni_model_type("gs://any-bucket/Fun-CosyVoice3-0.5B-2512") == "cosyvoice3" + class TestLoadAndResolveStageConfigs: def test_load_and_resolve_with_kwargs(self, mocker: MockerFixture): diff --git a/vllm_omni/config/config_factory.py b/vllm_omni/config/config_factory.py index a1778eda3..5adf218c7 100644 --- a/vllm_omni/config/config_factory.py +++ b/vllm_omni/config/config_factory.py @@ -15,6 +15,7 @@ from transformers import PretrainedConfig from vllm.logger import init_logger from vllm.transformers_utils.config import get_config from vllm.transformers_utils.repo_utils import get_hf_file_to_dict +from vllm.transformers_utils.runai_utils import ObjectStorageModel, is_runai_obj_uri from vllm_omni.config.endpoint_policy import EndpointRestriction from vllm_omni.config.omni_config import VllmOmniConfig @@ -45,6 +46,39 @@ logger = init_logger(__name__) _DEFAULT_PARALLEL_DEGREE = 1 +@functools.cache +def _materialize_object_storage_configs(model: str) -> str: + """Materialize an object-storage model URI's config files locally. + + vLLM's Run:AI streamer keeps ``s3://``/``gs://``/``az://`` URIs opaque until + each stage builds its ``ModelConfig``; parent-process resolution (HF config + lookup, pipeline/pipeline-key matching) would instead hand the URI to + ``huggingface_hub`` helpers, which reject it with ``HFValidationError``. + Pull the lightweight files once into vLLM's deterministic + ``model_streamer/`` directory so config reads work here, and so the + stage processes' own pull lands in that same directory. + + Returns the input unchanged for non object-storage paths. + """ + if not is_runai_obj_uri(model): + return model + object_storage_model = ObjectStorageModel(url=model) + object_storage_model.pull_files(model, allow_pattern=["*.model", "*.py", "*.json"]) + logger.info("Materialized object-storage configs for %s at %s", model, object_storage_model.dir) + return object_storage_model.dir + + +def _name_match_candidate(model: str) -> str: + """Last path component of a model reference, used for name-based matching. + + Object-storage URIs and HF repo ids carry non-model segments (bucket name, + organization) that must not participate in substring matching; e.g. a + bucket named ``qwen3-tts-models`` holding a ``Qwen3-Omni`` checkpoint must + not resolve to the ``qwen3_tts`` pipeline. + """ + return model.rstrip("/").rsplit("/", 1)[-1] + + def with_trust_remote_code_override( overrides: Mapping[str, Any], trust_remote_code: bool | None, @@ -117,7 +151,7 @@ class StageConfigFactory: """ hf_config = None try: - return get_config(model, trust_remote_code=trust_remote_code) + return get_config(_materialize_object_storage_configs(model), trust_remote_code=trust_remote_code) except Exception as e: logger.debug(f"`get_config` failed with exception {e}; inferred HF config is None") return hf_config @@ -163,10 +197,12 @@ class StageConfigFactory: if hf_config is not None: return hf_config.model_type + config_source = _materialize_object_storage_configs(model) + # Fallback: read config.json directly for custom model types that # are not registered with transformers (e.g. qwen3_tts). try: - config_dict = get_hf_file_to_dict("config.json", model, revision=None) + config_dict = get_hf_file_to_dict("config.json", config_source, revision=None) if config_dict: if "model_type" in config_dict: return config_dict["model_type"] @@ -183,7 +219,7 @@ class StageConfigFactory: # model_index.json with _class_name that maps to a pipeline key via # PipelineConfig.diffusers_class_name. try: - model_index = get_hf_file_to_dict("model_index.json", model, revision=None) + model_index = get_hf_file_to_dict("model_index.json", config_source, revision=None) if model_index and "_class_name" in model_index: class_name = model_index["_class_name"] for obj in OMNI_PIPELINES.values(): @@ -203,8 +239,10 @@ class StageConfigFactory: # Final fallback: some models (e.g. CosyVoice3) ship an empty # config.json and rely on naming conventions. Match the model path # basename against registered pipeline keys — longest match wins - # so "cosyvoice3" (length 10) beats "cosyvoice" (length 9). - model_lower = model.lower().replace("-", "").replace("_", "") + # so "cosyvoice3" (length 10) beats "cosyvoice" (length 9). Only + # the basename is scanned so URI segments such as the bucket name + # cannot select an unrelated pipeline. + model_lower = _name_match_candidate(model).lower().replace("-", "").replace("_", "") best: str | None = None best_len = 0 for registered_key in OMNI_PIPELINES.keys(): diff --git a/vllm_omni/entrypoints/omni_base.py b/vllm_omni/entrypoints/omni_base.py index b787c44a4..ab2bee01b 100644 --- a/vllm_omni/entrypoints/omni_base.py +++ b/vllm_omni/entrypoints/omni_base.py @@ -10,6 +10,7 @@ import huggingface_hub import vllm.envs as envs from vllm.logger import init_logger from vllm.transformers_utils.repo_utils import file_or_path_exists +from vllm.transformers_utils.runai_utils import is_runai_obj_uri from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError from vllm_omni.engine.async_omni_engine import AsyncOmniEngine @@ -66,6 +67,14 @@ def omni_snapshot_download(model_id: str) -> str: if os.path.exists(model_id): return model_id + # Object-storage models must remain URIs until each stage constructs its + # ModelConfig. vLLM then materializes only config/tokenizer files locally + # and keeps the URI in model_weights for Run:AI streaming. Treating the URI + # as a Hugging Face repo here either fails validation or downloads through + # the wrong backend before the stage processes are created. + if is_runai_obj_uri(model_id): + return model_id + # TODO: this is just a workaround for quickly use modelscope, we should support # modelscope in weight loading feature instead of using `snapshot_download` # Read through ``vllm.envs`` so the flag keeps vLLM's semantics (only diff --git a/vllm_omni/entrypoints/utils.py b/vllm_omni/entrypoints/utils.py index 07b45f374..9e884db48 100644 --- a/vllm_omni/entrypoints/utils.py +++ b/vllm_omni/entrypoints/utils.py @@ -11,7 +11,12 @@ from vllm.sampling_params import RequestOutputKind, SamplingParams 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.config_factory import ( + StageConfigFactory, + _materialize_object_storage_configs, + _name_match_candidate, + 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 @@ -210,12 +215,13 @@ def _try_resolve_omni_model_type(model: str) -> str | None: Searches both the legacy ``stage_configs/*.yaml`` directory and the migrated ``deploy/*.yaml`` directory for a stem that substring-matches - the model path (e.g. ``cosyvoice3`` in - ``FunAudioLLM/Fun-CosyVoice3-0.5B-2512``). The longest match wins so - ``cosyvoice3`` beats ``cosyvoice`` and ``bagel_single_stage`` beats - ``bagel``. + the model path basename (e.g. ``cosyvoice3`` in + ``FunAudioLLM/Fun-CosyVoice3-0.5B-2512``). Only the basename is scanned + so URI segments such as the bucket name cannot select an unrelated + pipeline. The longest match wins so ``cosyvoice3`` beats ``cosyvoice`` + and ``bagel_single_stage`` beats ``bagel``. """ - model_lower = model.lower().replace("-", "").replace("_", "") + model_lower = _name_match_candidate(model).lower().replace("-", "").replace("_", "") best_match: str | None = None best_len = 0 for subdir in ("model_executor/stage_configs", "deploy"): @@ -274,24 +280,29 @@ def resolve_model_config_path(model: str) -> str | None: Raises: ValueError: If model_type cannot be determined """ + # Object-storage URIs are streamed by vLLM's Run:AI streamer only after + # each stage builds its ModelConfig, so config resolution here reads the + # local copy that vLLM-Omni materializes for such URIs. Name-based + # fallbacks keep the original string (the materialized path is a hash). + config_source = _materialize_object_storage_configs(model) # Try to get config from standard transformers format first try: - hf_config = get_config(model, trust_remote_code=True) + hf_config = get_config(config_source, trust_remote_code=True) model_type = hf_config.model_type except (ValueError, Exception): # If standard transformers format fails, try diffusers format - if get_diffusion_model_index(model) is not None: - model_type = _try_get_class_name_from_diffusers_config(model) + if get_diffusion_model_index(config_source) is not None: + model_type = _try_get_class_name_from_diffusers_config(config_source) if model_type is None: raise ValueError( f"Could not determine model_type for diffusers model: {model}. " "Please ensure its Diffusers pipeline index contains '_class_name'" ) - elif file_or_path_exists(model, "config.json", revision=None): + elif file_or_path_exists(config_source, "config.json", revision=None): # Try to read config.json manually for custom models like Bagel that fail get_config # but have a valid config.json with model_type try: - config_dict = get_hf_file_to_dict("config.json", model, revision=None) + config_dict = get_hf_file_to_dict("config.json", config_source, revision=None) if config_dict and "model_type" in config_dict: model_type = config_dict["model_type"] else: @@ -315,7 +326,7 @@ def resolve_model_config_path(model: str) -> str | None: ) default_config_path = current_omni_platform.get_default_stage_config_path() - if model_type == "vla" and _looks_like_dreamzero(model): + if model_type == "vla" and _looks_like_dreamzero(config_source): model_type = "dreamzero" if model_type in _DIFFUSERS_CLASS_TO_CONFIG: