[Bugfix] Preserve object storage URIs during model resolution (#5036)

Signed-off-by: Ziyang Zhang <hafuhafu@qq.com>

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"

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.

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

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/<hash>`` 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():

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

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: