[Bugfix][MiniMax-Music3] Resolve stage subdirs against a real snapshot (#6640)

Signed-off-by: Yueqian Lin <linyueqian@outlook.com>

tests/engine/test_async_omni_engine_stage_init.py

@@ -4,6 +4,7 @@
 import concurrent.futures
 import contextlib
 import importlib
+import json
 import os
 import time
 import types
@@ -1093,6 +1094,8 @@ def test_build_engine_args_cli_tokenizer_overrides_inferred_base_tokenizer(tmp_p
         default_sampling_params={},
     )
 
+    (tmp_path / "llm").mkdir()
+
     engine_args = build_engine_args_dict(
         stage_cfg,
         str(tmp_path),
@@ -1154,6 +1157,9 @@ def test_build_engine_args_keeps_stage_owned_tokenizer_subdir(tmp_path):
         default_sampling_params={},
     )
 
+    (tmp_path / "llm").mkdir()
+    (tmp_path / "tokenizer").mkdir()
+
     engine_args = build_engine_args_dict(
         stage_cfg,
         str(tmp_path),
@@ -1164,6 +1170,430 @@ def test_build_engine_args_keeps_stage_owned_tokenizer_subdir(tmp_path):
     assert engine_args["tokenizer"] == os.path.join(str(tmp_path), "tokenizer")
 
 
+def test_build_engine_args_rejects_missing_subdir_in_local_model_dir(tmp_path):
+    """A local model directory without the stage subfolder must fail closed.
+
+    Forwarding the joined path as ``EngineArgs.model`` sends a non-directory
+    string to HuggingFace, which reports it as a malformed repo id.
+    """
+    from vllm_omni.engine.stage_init_utils import build_engine_args_dict
+
+    stage_cfg = types.SimpleNamespace(
+        stage_id=0,
+        stage_type="llm",
+        engine_args={"model_subdir": "language_model", "tokenizer_subdir": "tokenizer"},
+        default_sampling_params={},
+    )
+    (tmp_path / "tokenizer").mkdir()
+
+    with pytest.raises(RuntimeError, match="language_model"):
+        build_engine_args_dict(stage_cfg, str(tmp_path))
+
+
+# A real snapshot subfolder always carries artifacts; empty directories would
+# assert the broken exists-means-complete predicate this suite regresses.
+_SUBDIR_ARTIFACT = {"language_model": "model.safetensors", "tokenizer": "tokenizer.json"}
+
+
+def _make_snapshot(root, subdirs):
+    for subdir in subdirs:
+        folder = root / subdir
+        folder.mkdir(parents=True, exist_ok=True)
+        (folder / _SUBDIR_ARTIFACT.get(subdir, "data.bin")).write_text("x")
+    return str(root)
+
+
+def test_build_engine_args_pulls_stage_subdirs_missing_from_cached_snapshot(monkeypatch, tmp_path):
+    """Regression for issue #6638.
+
+    ``snapshot_download(local_files_only=True)`` hands back a snapshot root as
+    soon as the repo is partly cached, even when this stage's subfolders were
+    never materialized. Stage init must fetch exactly those subfolders instead
+    of joining onto a path that exists nowhere.
+    """
+    from huggingface_hub import HfApi
+
+    from vllm_omni.engine.stage_init_utils import build_engine_args_dict
+
+    snapshot = tmp_path / "snapshots" / "deadbeef"
+    _make_snapshot(snapshot, ["tokenizer"])
+    calls = []
+
+    def fake_snapshot_download(self, repo_id, **kwargs):
+        calls.append(kwargs)
+        if kwargs.get("local_files_only"):
+            return str(snapshot)
+        _make_snapshot(snapshot, ["language_model"])
+        return str(snapshot)
+
+    monkeypatch.setattr(HfApi, "snapshot_download", fake_snapshot_download)
+
+    stage_cfg = types.SimpleNamespace(
+        stage_id=0,
+        stage_type="llm",
+        engine_args={"model_subdir": "language_model", "tokenizer_subdir": "tokenizer"},
+        default_sampling_params={},
+    )
+
+    engine_args = build_engine_args_dict(stage_cfg, "MiniMaxAI/MiniMax-Music3")
+
+    assert engine_args["model"] == os.path.join(str(snapshot), "language_model")
+    assert engine_args["tokenizer"] == os.path.join(str(snapshot), "tokenizer")
+    assert os.path.isdir(engine_args["model"])
+    # Only the stage's own subfolders are pulled, not the whole repo.
+    assert calls[-1]["allow_patterns"] == ["language_model/*", "tokenizer/*"]
+
+
+def test_build_engine_args_skips_hub_call_when_cached_snapshot_is_complete(monkeypatch, tmp_path):
+    """A warm cache stays offline-friendly: no outgoing Hub request."""
+    from huggingface_hub import HfApi
+
+    from vllm_omni.engine.stage_init_utils import build_engine_args_dict
+
+    snapshot = tmp_path / "snapshots" / "deadbeef"
+    _make_snapshot(snapshot, ["language_model", "tokenizer"])
+
+    def fake_snapshot_download(self, repo_id, **kwargs):
+        assert kwargs.get("local_files_only"), "warm cache must not reach the Hub"
+        return str(snapshot)
+
+    monkeypatch.setattr(HfApi, "snapshot_download", fake_snapshot_download)
+
+    stage_cfg = types.SimpleNamespace(
+        stage_id=0,
+        stage_type="llm",
+        engine_args={"model_subdir": "language_model", "tokenizer_subdir": "tokenizer"},
+        default_sampling_params={},
+    )
+
+    engine_args = build_engine_args_dict(stage_cfg, "MiniMaxAI/MiniMax-Music3")
+
+    assert engine_args["model"] == os.path.join(str(snapshot), "language_model")
+
+
+def test_build_engine_args_redownloads_a_partial_subdir(monkeypatch, tmp_path):
+    """A subfolder holding only config.json is an interrupted download.
+
+    The warm-cache early return must not accept it: converting the Hub ID into
+    that local path leaves vLLM's loader without any fallback for the missing
+    weights.
+    """
+    from huggingface_hub import HfApi
+
+    from vllm_omni.engine.stage_init_utils import build_engine_args_dict
+
+    snapshot = tmp_path / "snapshots" / "deadbeef"
+    partial = snapshot / "language_model"
+    partial.mkdir(parents=True)
+    (partial / "config.json").write_text("{}")
+    _make_snapshot(snapshot, ["tokenizer"])
+    calls = []
+
+    def fake_snapshot_download(self, repo_id, **kwargs):
+        calls.append(kwargs)
+        if kwargs.get("local_files_only"):
+            return str(snapshot)
+        _make_snapshot(snapshot, ["language_model"])
+        return str(snapshot)
+
+    monkeypatch.setattr(HfApi, "snapshot_download", fake_snapshot_download)
+
+    stage_cfg = types.SimpleNamespace(
+        stage_id=0,
+        stage_type="llm",
+        engine_args={"model_subdir": "language_model", "tokenizer_subdir": "tokenizer"},
+        default_sampling_params={},
+    )
+
+    engine_args = build_engine_args_dict(stage_cfg, "MiniMaxAI/MiniMax-Music3")
+
+    assert any(not call.get("local_files_only") for call in calls), "a partial subdir must fall through to the Hub"
+    assert engine_args["model"] == os.path.join(str(snapshot), "language_model")
+
+
+def test_build_engine_args_redownloads_when_index_lists_missing_shards(monkeypatch, tmp_path):
+    """An index plus a subset of its shards is an interrupted download.
+
+    The index downloads early; accepting it as proof of weights converts the
+    Hub ID into a local path vLLM cannot fetch the remaining shards for.
+    """
+    from huggingface_hub import HfApi
+
+    from vllm_omni.engine.stage_init_utils import build_engine_args_dict
+
+    snapshot = tmp_path / "snapshots" / "deadbeef"
+    partial = snapshot / "language_model"
+    partial.mkdir(parents=True)
+    (partial / "model.safetensors.index.json").write_text(
+        json.dumps(
+            {
+                "weight_map": {
+                    "a.weight": "model-00001-of-00002.safetensors",
+                    "b.weight": "model-00002-of-00002.safetensors",
+                }
+            }
+        )
+    )
+    (partial / "model-00001-of-00002.safetensors").write_text("x")
+    _make_snapshot(snapshot, ["tokenizer"])
+    calls = []
+
+    def fake_snapshot_download(self, repo_id, **kwargs):
+        calls.append(kwargs)
+        if not kwargs.get("local_files_only"):
+            (partial / "model-00002-of-00002.safetensors").write_text("x")
+        return str(snapshot)
+
+    monkeypatch.setattr(HfApi, "snapshot_download", fake_snapshot_download)
+
+    stage_cfg = types.SimpleNamespace(
+        stage_id=0,
+        stage_type="llm",
+        engine_args={"model_subdir": "language_model", "tokenizer_subdir": "tokenizer"},
+        default_sampling_params={},
+    )
+
+    engine_args = build_engine_args_dict(stage_cfg, "MiniMaxAI/MiniMax-Music3")
+
+    assert any(not call.get("local_files_only") for call in calls), "missing shards must fall through to the Hub"
+    assert engine_args["model"] == os.path.join(str(snapshot), "language_model")
+
+
+def test_build_engine_args_redownloads_shards_without_their_index(monkeypatch, tmp_path):
+    """Shard-named weights always ship an index; one without it is partial."""
+    from huggingface_hub import HfApi
+
+    from vllm_omni.engine.stage_init_utils import build_engine_args_dict
+
+    snapshot = tmp_path / "snapshots" / "deadbeef"
+    partial = snapshot / "language_model"
+    partial.mkdir(parents=True)
+    (partial / "model-00001-of-00004.safetensors").write_text("x")
+    _make_snapshot(snapshot, ["tokenizer"])
+    calls = []
+
+    def fake_snapshot_download(self, repo_id, **kwargs):
+        calls.append(kwargs)
+        if not kwargs.get("local_files_only"):
+            shards = {f"w{i}.weight": f"model-0000{i}-of-00004.safetensors" for i in range(1, 5)}
+            (partial / "model.safetensors.index.json").write_text(json.dumps({"weight_map": shards}))
+            for shard in shards.values():
+                (partial / shard).write_text("x")
+        return str(snapshot)
+
+    monkeypatch.setattr(HfApi, "snapshot_download", fake_snapshot_download)
+
+    stage_cfg = types.SimpleNamespace(
+        stage_id=0,
+        stage_type="llm",
+        engine_args={"model_subdir": "language_model", "tokenizer_subdir": "tokenizer"},
+        default_sampling_params={},
+    )
+
+    build_engine_args_dict(stage_cfg, "MiniMaxAI/MiniMax-Music3")
+
+    assert any(not call.get("local_files_only") for call in calls)
+
+
+def test_build_engine_args_redownloads_a_tokenizer_folder_without_vocabulary(monkeypatch, tmp_path):
+    """Templates and configs download first; alone they are not a tokenizer."""
+    from huggingface_hub import HfApi
+
+    from vllm_omni.engine.stage_init_utils import build_engine_args_dict
+
+    snapshot = tmp_path / "snapshots" / "deadbeef"
+    _make_snapshot(snapshot, ["language_model"])
+    partial = snapshot / "tokenizer"
+    partial.mkdir(parents=True)
+    (partial / "chat_template.jinja").write_text("{{ messages }}")
+    calls = []
+
+    def fake_snapshot_download(self, repo_id, **kwargs):
+        calls.append(kwargs)
+        if not kwargs.get("local_files_only"):
+            _make_snapshot(snapshot, ["tokenizer"])
+        return str(snapshot)
+
+    monkeypatch.setattr(HfApi, "snapshot_download", fake_snapshot_download)
+
+    stage_cfg = types.SimpleNamespace(
+        stage_id=0,
+        stage_type="llm",
+        engine_args={"model_subdir": "language_model", "tokenizer_subdir": "tokenizer"},
+        default_sampling_params={},
+    )
+
+    engine_args = build_engine_args_dict(stage_cfg, "MiniMaxAI/MiniMax-Music3")
+
+    assert any(not call.get("local_files_only") for call in calls)
+    assert engine_args["tokenizer"] == os.path.join(str(snapshot), "tokenizer")
+
+
+def test_build_engine_args_forwards_revision_and_download_dir(monkeypatch, tmp_path):
+    """revision/download_dir must shape snapshot selection (they cannot be
+    corrected downstream once the repo ID is a local path)."""
+    from huggingface_hub import HfApi
+
+    from vllm_omni.engine.stage_init_utils import build_engine_args_dict
+
+    snapshot = tmp_path / "snapshots" / "pinned"
+    _make_snapshot(snapshot, ["language_model", "tokenizer"])
+    seen = []
+
+    def fake_snapshot_download(self, repo_id, **kwargs):
+        seen.append(kwargs)
+        return str(snapshot)
+
+    monkeypatch.setattr(HfApi, "snapshot_download", fake_snapshot_download)
+
+    stage_cfg = types.SimpleNamespace(
+        stage_id=0,
+        stage_type="llm",
+        engine_args={
+            "model_subdir": "language_model",
+            "tokenizer_subdir": "tokenizer",
+            "revision": "v2.0",
+            "download_dir": str(tmp_path / "cache"),
+        },
+        default_sampling_params={},
+    )
+
+    engine_args = build_engine_args_dict(stage_cfg, "MiniMaxAI/MiniMax-Music3")
+
+    assert engine_args["model"] == os.path.join(str(snapshot), "language_model")
+    assert seen, "snapshot resolution must have been consulted"
+    for call in seen:
+        assert call.get("revision") == "v2.0"
+        assert call.get("cache_dir") == str(tmp_path / "cache")
+
+
+def test_build_engine_args_resolves_tokenizer_revision_separately(monkeypatch, tmp_path):
+    """A tokenizer pinned to a different revision resolves against its own snapshot."""
+    from huggingface_hub import HfApi
+
+    from vllm_omni.engine.stage_init_utils import build_engine_args_dict
+
+    model_snapshot = tmp_path / "snapshots" / "model-rev"
+    _make_snapshot(model_snapshot, ["language_model"])
+    tokenizer_snapshot = tmp_path / "snapshots" / "tok-rev"
+    _make_snapshot(tokenizer_snapshot, ["tokenizer"])
+
+    def fake_snapshot_download(self, repo_id, **kwargs):
+        if kwargs.get("revision") == "tok-rev":
+            return str(tokenizer_snapshot)
+        return str(model_snapshot)
+
+    monkeypatch.setattr(HfApi, "snapshot_download", fake_snapshot_download)
+
+    stage_cfg = types.SimpleNamespace(
+        stage_id=0,
+        stage_type="llm",
+        engine_args={
+            "model_subdir": "language_model",
+            "tokenizer_subdir": "tokenizer",
+            "revision": "model-rev",
+            "tokenizer_revision": "tok-rev",
+        },
+        default_sampling_params={},
+    )
+
+    engine_args = build_engine_args_dict(stage_cfg, "MiniMaxAI/MiniMax-Music3")
+
+    assert engine_args["model"] == os.path.join(str(model_snapshot), "language_model")
+    assert engine_args["tokenizer"] == os.path.join(str(tokenizer_snapshot), "tokenizer")
+
+
+def test_build_engine_args_resolves_root_tokenizer_revision_separately(monkeypatch, tmp_path):
+    """An empty tokenizer_subdir targets the snapshot root and still honors
+    its own revision."""
+    from huggingface_hub import HfApi
+
+    from vllm_omni.engine.stage_init_utils import build_engine_args_dict
+
+    model_snapshot = tmp_path / "snapshots" / "model-rev"
+    _make_snapshot(model_snapshot, ["language_model"])
+    # An empty tokenizer_subdir means the tokenizer IS the snapshot root, so
+    # the vocabulary artifact must sit at the root. Building it under
+    # ``tokenizer/`` returned a root that holds no loadable tokenizer, which
+    # the root-tokenizer validation now correctly rejects.
+    tokenizer_snapshot = tmp_path / "snapshots" / "tok-rev"
+    tokenizer_snapshot.mkdir(parents=True, exist_ok=True)
+    (tokenizer_snapshot / _SUBDIR_ARTIFACT["tokenizer"]).write_text("x")
+
+    def fake_snapshot_download(self, repo_id, **kwargs):
+        if kwargs.get("revision") == "tok-rev":
+            return str(tokenizer_snapshot)
+        return str(model_snapshot)
+
+    monkeypatch.setattr(HfApi, "snapshot_download", fake_snapshot_download)
+
+    stage_cfg = types.SimpleNamespace(
+        stage_id=0,
+        stage_type="llm",
+        engine_args={
+            "model_subdir": "language_model",
+            "tokenizer_subdir": "",
+            "revision": "model-rev",
+            "tokenizer_revision": "tok-rev",
+        },
+        default_sampling_params={},
+    )
+
+    engine_args = build_engine_args_dict(stage_cfg, "MiniMaxAI/MiniMax-Music3")
+
+    assert engine_args["model"] == os.path.join(str(model_snapshot), "language_model")
+    assert engine_args["tokenizer"] == str(tokenizer_snapshot)
+
+
+def test_build_engine_args_fails_closed_when_subdir_cannot_be_downloaded(monkeypatch, tmp_path):
+    """An undownloadable subfolder raises here instead of reaching HuggingFace."""
+    from huggingface_hub import HfApi
+
+    from vllm_omni.engine.stage_init_utils import build_engine_args_dict
+
+    snapshot = tmp_path / "snapshots" / "deadbeef"
+    _make_snapshot(snapshot, ["tokenizer"])
+
+    def fake_snapshot_download(self, repo_id, **kwargs):
+        if kwargs.get("local_files_only"):
+            return str(snapshot)
+        raise OSError("offline")
+
+    monkeypatch.setattr(HfApi, "snapshot_download", fake_snapshot_download)
+
+    stage_cfg = types.SimpleNamespace(
+        stage_id=0,
+        stage_type="llm",
+        engine_args={"model_subdir": "language_model", "tokenizer_subdir": "tokenizer"},
+        default_sampling_params={},
+    )
+
+    with pytest.raises(RuntimeError, match="language_model"):
+        build_engine_args_dict(stage_cfg, "MiniMaxAI/MiniMax-Music3")
+
+
+def test_build_engine_args_fails_closed_on_cold_cache_instead_of_joining_repo_id(monkeypatch):
+    """With nothing cached, the stage must not join a subdir onto the repo id."""
+    from huggingface_hub import HfApi
+
+    from vllm_omni.engine.stage_init_utils import build_engine_args_dict
+
+    def fake_snapshot_download(self, repo_id, **kwargs):
+        raise OSError("offline")
+
+    monkeypatch.setattr(HfApi, "snapshot_download", fake_snapshot_download)
+
+    stage_cfg = types.SimpleNamespace(
+        stage_id=0,
+        stage_type="llm",
+        engine_args={"model_subdir": "language_model", "tokenizer_subdir": "tokenizer"},
+        default_sampling_params={},
+    )
+
+    with pytest.raises(RuntimeError, match="nothing is cached locally"):
+        build_engine_args_dict(stage_cfg, "MiniMaxAI/MiniMax-Music3")
+
+
 def test_model_path_resolver_is_generic_and_model_owned(tmp_path):
     from vllm_omni.engine.stage_init_utils import _resolve_model_path
 

tests/engine/test_stage_engine_args.py

@@ -630,8 +630,24 @@ def test_typed_ming_image_engine_args_defer_diffusion_batch_default():
     assert OmniDiffusionConfig(**typed_backend_args).max_num_seqs == 1
 
 
+def _fake_model_root(pipeline, tmp_path):
+    """A model directory carrying every subfolder the pipeline's stages declare.
+
+    Stage init fails closed when a declared ``model_subdir``/``tokenizer_subdir``
+    is not a real directory, because the joined path would otherwise reach
+    HuggingFace as a malformed repo id (issue #6638).
+    """
+    root = tmp_path / "model"
+    root.mkdir(exist_ok=True)
+    for stage in pipeline.stages:
+        for subdir in (stage.model_subdir, stage.tokenizer_subdir):
+            if subdir:
+                (root / subdir).mkdir(parents=True, exist_ok=True)
+    return str(root)
+
+
 @pytest.mark.parametrize("model_type", sorted(OMNI_PIPELINES))
-def test_typed_engine_args_match_current_registry_backend_semantics(model_type):
+def test_typed_engine_args_match_current_registry_backend_semantics(model_type, tmp_path):
     pipeline = resolve_pipeline_config(model_type)
     if pipeline is None:
         pytest.skip(f"Pipeline {model_type!r} requires an HF config to resolve")
@@ -641,18 +657,19 @@ def test_typed_engine_args_match_current_registry_backend_semantics(model_type):
         if pipeline.default_deploy_config_name is not None
         else DeployConfig()
     )
+    model = _fake_model_root(pipeline, tmp_path)
     legacy_stages, omni_config = _legacy_and_typed_stages(
         pipeline,
         deploy,
-        model="/tmp",
+        model=model,
     )
 
     for legacy_stage in legacy_stages:
         stage_id = legacy_stage.stage_id
-        legacy_args = build_legacy_engine_args_dict(legacy_stage, model="/tmp")
+        legacy_args = build_legacy_engine_args_dict(legacy_stage, model=model)
         typed_args = build_engine_args_dict_from_omni_stage_config(
             omni_config.stage_by_id(stage_id),
-            model="/tmp",
+            model=model,
         )
         if legacy_stage.stage_type == StageType.DIFFUSION:
             backend_fields = _DIFFUSION_BACKEND_FIELDS

tests/model_executor/models/test_minimax_music3_repo_root.py

@@ -0,0 +1,339 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Repository-root resolution for MiniMax Music 3's multi-component checkpoint.
+
+The acoustic stage declares no ``model_subdir`` because it reads the repo root
+itself, so on a hub-id deployment its model path is ``MiniMaxAI/MiniMax-Music3``
+rather than a directory. Resolving that as a relative path silently points at
+the server's working directory (issue #6638).
+"""
+
+from __future__ import annotations
+
+import pytest
+
+from vllm_omni.model_executor.models.minimax_music3.weights import (
+    _COMPONENT_WEIGHT_NAME,
+    _REQUIRED_COMPONENTS,
+    _ROOT_MARKERS,
+    resolve_repo_root,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _make_root(path):
+    # A real snapshot carries weights for every component a stage loads, not
+    # just the folders that identify the root. Building only ``_ROOT_MARKERS``
+    # here would let a snapshot missing ``condition_encoder`` pass as complete,
+    # which is the gap this suite regresses.
+    for component in _REQUIRED_COMPONENTS:
+        folder = path / component
+        folder.mkdir(parents=True, exist_ok=True)
+        (folder / _COMPONENT_WEIGHT_NAME).write_bytes(b"x")
+    return path
+
+
+def test_resolve_repo_root_accepts_the_root_itself(tmp_path):
+    root = _make_root(tmp_path / "snapshot")
+    assert resolve_repo_root(str(root)) == root
+
+
+def test_resolve_repo_root_walks_up_from_a_model_subdir(tmp_path):
+    root = _make_root(tmp_path / "snapshot")
+    (root / "language_model").mkdir()
+    assert resolve_repo_root(str(root / "language_model")) == root
+
+
+def test_resolve_repo_root_resolves_a_hub_id_to_a_local_snapshot(monkeypatch, tmp_path):
+    """A hub id must resolve through the cache, not against the working directory."""
+    from huggingface_hub import HfApi
+
+    root = _make_root(tmp_path / "snapshots" / "deadbeef")
+    seen = []
+
+    def fake_snapshot_download(self, repo_id, **kwargs):
+        seen.append(kwargs)
+        return str(root)
+
+    monkeypatch.setattr(HfApi, "snapshot_download", fake_snapshot_download)
+
+    assert resolve_repo_root("MiniMaxAI/MiniMax-Music3") == root
+    # Cache-first, and only the component folders.
+    assert seen[0]["local_files_only"] is True
+    assert "transformer/*" in seen[0]["allow_patterns"]
+    assert "qwen_7B/*" not in seen[0]["allow_patterns"]
+
+
+def test_resolve_repo_root_falls_back_to_the_hub_when_the_cache_is_incomplete(monkeypatch, tmp_path):
+    from huggingface_hub import HfApi
+
+    root = _make_root(tmp_path / "snapshots" / "deadbeef")
+    calls = []
+
+    def fake_snapshot_download(self, repo_id, **kwargs):
+        calls.append(kwargs)
+        if kwargs.get("local_files_only"):
+            raise OSError("incomplete snapshot")
+        return str(root)
+
+    monkeypatch.setattr(HfApi, "snapshot_download", fake_snapshot_download)
+
+    assert resolve_repo_root("MiniMaxAI/MiniMax-Music3") == root
+    assert len(calls) == 2
+
+
+def test_resolve_repo_root_retries_online_when_the_local_lookup_returns_a_partial_root(monkeypatch, tmp_path):
+    """Hub clients differ on partial caches: some raise, some return the root.
+
+    A returned partial root must fall through to the online call all the same,
+    or startup fails without ever consulting the Hub.
+    """
+    from huggingface_hub import HfApi
+
+    snapshot = tmp_path / "models--MiniMaxAI--MiniMax-Music3" / "snapshots" / "deadbeef"
+    (snapshot / "language_model").mkdir(parents=True)
+    calls = []
+
+    def fake_snapshot_download(self, repo_id, **kwargs):
+        calls.append(kwargs)
+        if not kwargs.get("local_files_only"):
+            _make_root(snapshot)
+        return str(snapshot)
+
+    monkeypatch.setattr(HfApi, "snapshot_download", fake_snapshot_download)
+
+    assert resolve_repo_root("MiniMaxAI/MiniMax-Music3") == snapshot
+    assert [bool(call.get("local_files_only")) for call in calls] == [True, False]
+
+
+def test_resolve_repo_root_redownloads_weightless_marker_dirs(monkeypatch, tmp_path):
+    """Marker folders without their component weights are an interrupted download."""
+    from huggingface_hub import HfApi
+
+    snapshot = tmp_path / "models--MiniMaxAI--MiniMax-Music3" / "snapshots" / "deadbeef"
+    for marker in _ROOT_MARKERS:
+        (snapshot / marker).mkdir(parents=True)
+    calls = []
+
+    def fake_snapshot_download(self, repo_id, **kwargs):
+        calls.append(kwargs)
+        if not kwargs.get("local_files_only"):
+            _make_root(snapshot)
+        return str(snapshot)
+
+    monkeypatch.setattr(HfApi, "snapshot_download", fake_snapshot_download)
+
+    assert resolve_repo_root("MiniMaxAI/MiniMax-Music3") == snapshot
+    assert any(not call.get("local_files_only") for call in calls)
+
+
+def test_resolve_repo_root_reports_the_original_reference_when_unresolvable(monkeypatch):
+    from huggingface_hub import HfApi
+
+    def fake_snapshot_download(self, repo_id, **kwargs):
+        raise OSError("offline")
+
+    monkeypatch.setattr(HfApi, "snapshot_download", fake_snapshot_download)
+
+    with pytest.raises(FileNotFoundError, match="MiniMaxAI/MiniMax-Music3"):
+        resolve_repo_root("MiniMaxAI/MiniMax-Music3")
+
+
+def test_resolve_repo_root_downloads_components_for_a_cold_cache_model_subdir(monkeypatch, tmp_path):
+    """Cold-cache stage 0: the talker gets an EXISTING ``language_model`` dir.
+
+    Stage init pre-downloads only ``language_model/`` and ``tokenizer/``, so
+    the marker walk fails although the repo id is recoverable from the cache
+    layout and one snapshot_download away from working.
+    """
+    from huggingface_hub import HfApi
+
+    snapshot = tmp_path / "models--MiniMaxAI--MiniMax-Music3" / "snapshots" / "deadbeef"
+    (snapshot / "language_model").mkdir(parents=True)
+    (snapshot / "tokenizer").mkdir()
+    seen = []
+
+    def fake_snapshot_download(self, repo_id, **kwargs):
+        seen.append(repo_id)
+        _make_root(snapshot)
+        return str(snapshot)
+
+    monkeypatch.setattr(HfApi, "snapshot_download", fake_snapshot_download)
+
+    assert resolve_repo_root(str(snapshot / "language_model")) == snapshot
+    assert seen == ["MiniMaxAI/MiniMax-Music3"]
+
+
+def test_resolve_repo_root_does_not_guess_a_repo_id_for_plain_directories(monkeypatch, tmp_path):
+    """A directory outside the HF cache has no recoverable repo id."""
+    from huggingface_hub import HfApi
+
+    def fake_snapshot_download(self, repo_id, **kwargs):
+        raise AssertionError("must not reach the Hub for a plain directory")
+
+    monkeypatch.setattr(HfApi, "snapshot_download", fake_snapshot_download)
+
+    plain = tmp_path / "language_model"
+    plain.mkdir()
+    with pytest.raises(FileNotFoundError, match="expected sibling folders"):
+        resolve_repo_root(str(plain))
+
+
+def test_resolve_repo_root_rejects_a_snapshot_missing_condition_encoder(monkeypatch, tmp_path):
+    """``condition_encoder`` is loaded by the acoustic stage but is not a root marker.
+
+    A cache holding only the marker folders satisfies the name probe while the
+    acoustic stage would still fail, so completeness must cover every required
+    component rather than the identification markers alone.
+    """
+    from huggingface_hub import HfApi
+
+    snapshot = tmp_path / "models--MiniMaxAI--MiniMax-Music3" / "snapshots" / "deadbeef"
+    for marker in _ROOT_MARKERS:
+        folder = snapshot / marker
+        folder.mkdir(parents=True)
+        (folder / _COMPONENT_WEIGHT_NAME).write_bytes(b"x")
+    repaired = []
+
+    def fake_snapshot_download(self, repo_id, **kwargs):
+        if not kwargs.get("local_files_only"):
+            repaired.append(repo_id)
+            _make_root(snapshot)
+        return str(snapshot)
+
+    monkeypatch.setattr(HfApi, "snapshot_download", fake_snapshot_download)
+
+    assert resolve_repo_root(str(snapshot)) == snapshot
+    assert repaired == ["MiniMaxAI/MiniMax-Music3"]
+
+
+def test_resolve_repo_root_rejects_a_partial_shard_run(monkeypatch, tmp_path):
+    """One shard of a multi-shard component satisfies a bare glob but cannot load."""
+    from huggingface_hub import HfApi
+
+    stem = _COMPONENT_WEIGHT_NAME.removesuffix(".safetensors")
+    snapshot = tmp_path / "models--MiniMaxAI--MiniMax-Music3" / "snapshots" / "deadbeef"
+    for component in _REQUIRED_COMPONENTS:
+        (snapshot / component).mkdir(parents=True)
+        (snapshot / component / _COMPONENT_WEIGHT_NAME).write_bytes(b"x")
+    # Replace one component with an interrupted 1-of-3 shard run.
+    partial = snapshot / _ROOT_MARKERS[0]
+    (partial / _COMPONENT_WEIGHT_NAME).unlink()
+    (partial / f"{stem}-00001-of-00003.safetensors").write_bytes(b"x")
+    repaired = []
+
+    def fake_snapshot_download(self, repo_id, **kwargs):
+        if not kwargs.get("local_files_only"):
+            repaired.append(repo_id)
+            for index in (2, 3):
+                (partial / f"{stem}-0000{index}-of-00003.safetensors").write_bytes(b"x")
+        return str(snapshot)
+
+    monkeypatch.setattr(HfApi, "snapshot_download", fake_snapshot_download)
+
+    assert resolve_repo_root(str(snapshot)) == snapshot
+    assert repaired == ["MiniMaxAI/MiniMax-Music3"]
+
+
+def test_resolve_repo_root_threads_revision_and_download_dir(monkeypatch, tmp_path):
+    """A pinned revision must not be dropped when fetching the components.
+
+    Resolving components against the default branch while the AR backbone is
+    pinned elsewhere would combine weights from two different commits.
+    """
+    from huggingface_hub import HfApi
+
+    snapshot = tmp_path / "snapshot"
+    calls = []
+
+    def fake_snapshot_download(self, repo_id, **kwargs):
+        calls.append(kwargs)
+        _make_root(snapshot)
+        return str(snapshot)
+
+    monkeypatch.setattr(HfApi, "snapshot_download", fake_snapshot_download)
+
+    assert resolve_repo_root("MiniMaxAI/MiniMax-Music3", revision="abc123", download_dir="/tmp/hub-cache") == snapshot
+    assert calls
+    assert all(call.get("revision") == "abc123" for call in calls)
+    assert all(call.get("cache_dir") == "/tmp/hub-cache" for call in calls)
+
+
+def test_resolve_repo_root_reuses_the_snapshot_revision_from_the_cache_path(monkeypatch, tmp_path):
+    """The commit is recoverable from ``snapshots/<sha>`` and must be reused.
+
+    Re-resolving the repo id alone would fetch the default branch rather than
+    the snapshot the caller is already sitting in.
+    """
+    from huggingface_hub import HfApi
+
+    snapshot = tmp_path / "models--MiniMaxAI--MiniMax-Music3" / "snapshots" / "cafebabe"
+    (snapshot / "language_model").mkdir(parents=True)
+    calls = []
+
+    def fake_snapshot_download(self, repo_id, **kwargs):
+        calls.append(kwargs)
+        _make_root(snapshot)
+        return str(snapshot)
+
+    monkeypatch.setattr(HfApi, "snapshot_download", fake_snapshot_download)
+
+    assert resolve_repo_root(str(snapshot / "language_model")) == snapshot
+    assert calls
+    assert all(call.get("revision") == "cafebabe" for call in calls)
+
+
+def test_resolve_repo_root_accepts_shard_names_the_loader_can_read(monkeypatch, tmp_path):
+    """The completeness predicate must not be stricter than ``load_component_state``.
+
+    The loader globs ``<stem>-*.safetensors`` and reads whatever it finds, so a
+    component whose weights do not use the ``-<index>-of-<total>`` naming is
+    loadable. Judging it incomplete would force a needless re-download and, with
+    no network, fail on weights that are present and readable.
+    """
+    from huggingface_hub import HfApi
+
+    stem = _COMPONENT_WEIGHT_NAME.removesuffix(".safetensors")
+    root = tmp_path / "snapshot"
+    for component in _REQUIRED_COMPONENTS:
+        (root / component).mkdir(parents=True)
+        (root / component / _COMPONENT_WEIGHT_NAME).write_bytes(b"x")
+    unsharded = root / _ROOT_MARKERS[0]
+    (unsharded / _COMPONENT_WEIGHT_NAME).unlink()
+    (unsharded / f"{stem}-fp16.safetensors").write_bytes(b"x")
+
+    def fake_snapshot_download(self, repo_id, **kwargs):
+        raise AssertionError("must not reach the Hub for readable weights")
+
+    monkeypatch.setattr(HfApi, "snapshot_download", fake_snapshot_download)
+
+    assert resolve_repo_root(str(root)) == root
+
+
+def test_resolve_repo_root_still_rejects_a_partial_run_beside_an_extra_file(monkeypatch, tmp_path):
+    """An unparsable extra must not mask a genuinely short ``-of-`` run."""
+    from huggingface_hub import HfApi
+
+    stem = _COMPONENT_WEIGHT_NAME.removesuffix(".safetensors")
+    root = tmp_path / "snapshot"
+    for component in _REQUIRED_COMPONENTS:
+        (root / component).mkdir(parents=True)
+        (root / component / _COMPONENT_WEIGHT_NAME).write_bytes(b"x")
+    partial = root / _ROOT_MARKERS[0]
+    (partial / _COMPONENT_WEIGHT_NAME).unlink()
+    (partial / f"{stem}-00001-of-00003.safetensors").write_bytes(b"x")
+    (partial / f"{stem}-extra.safetensors").write_bytes(b"x")
+    repaired = []
+
+    def fake_snapshot_download(self, repo_id, **kwargs):
+        repaired.append(repo_id)
+        for index in (2, 3):
+            (partial / f"{stem}-0000{index}-of-00003.safetensors").write_bytes(b"x")
+        return str(root)
+
+    monkeypatch.setattr(HfApi, "snapshot_download", fake_snapshot_download)
+
+    with pytest.raises(FileNotFoundError, match="components are incomplete"):
+        resolve_repo_root(str(root))
+    assert repaired == []

tools/pre_commit/check_forbidden_imports.py

@@ -194,7 +194,6 @@ CHECK_IMPORTS = {
             "vllm_omni/diffusion/models/utils.py",
             "vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py",
             "vllm_omni/engine/arg_utils.py",
-            "vllm_omni/engine/stage_init_utils.py",
             "vllm_omni/entrypoints/openai/api_server.py",
             "vllm_omni/entrypoints/openai/serving_speech.py",
             "vllm_omni/entrypoints/openai/tts_adapters/cosyvoice3.py",

vllm_omni/engine/stage_init_utils.py

@@ -14,19 +14,24 @@ from __future__ import annotations
 import copy
 import fcntl
 import importlib
+import json
 import multiprocessing as mp
 import os
 import time
-from collections.abc import Callable, Generator, Mapping, Sequence
+from collections.abc import Callable, Collection, Generator, Mapping, Sequence
 from contextlib import contextmanager
 from dataclasses import dataclass, fields, replace
+from pathlib import Path
 from typing import Any, Literal, cast
 
+import regex as re
 from vllm.logger import init_logger
 from vllm.pooling_params import PoolingParams
 from vllm.renderers import BaseRenderer
 from vllm.sampling_params import SamplingParams
 from vllm.tokenizers import cached_tokenizer_from_config
+from vllm.transformers_utils.repo_utils import hf_api
+from vllm.transformers_utils.runai_utils import is_runai_obj_uri
 from vllm.usage.usage_lib import UsageContext
 from vllm.v1.engine.input_processor import InputProcessor
 from vllm.v1.executor import Executor
@@ -82,22 +87,144 @@ class LogicalStageInitPlan:
     replicas: list[ReplicaInitPlan]
 
 
-def _resolve_model_to_local_path(model: str) -> str:
-    """Resolve an HF Hub model ID to a local cache path."""
+def _missing_stage_subdirs(base: str, subdirs: Sequence[str]) -> list[str]:
+    """Return the entries of ``subdirs`` that are not directories under ``base``."""
+    return [subdir for subdir in subdirs if not os.path.isdir(os.path.join(base, subdir))]
+
+
+# Artifacts that make a snapshot subfolder trustworthy. A directory that
+# exists but holds none of these is an interrupted download, not a snapshot;
+# treating it as complete strips the Hub fallback vLLM would need later.
+_WEIGHT_ARTIFACT_PATTERNS = ("*.safetensors", "*.bin", "*.pt", "*.gguf")
+# Vocabulary-bearing files. Configs and chat templates are small and download
+# first, so their presence alone cannot distinguish a tokenizer folder from an
+# interrupted download.
+_TOKENIZER_ARTIFACT_NAMES = (
+    "tokenizer.json",
+    "tokenizer.model",
+    "spiece.model",
+    "sentencepiece.bpe.model",
+    "vocab.json",
+    "vocab.txt",
+)
+# HF sharded checkpoints name their pieces `<stem>-NNNNN-of-NNNNN.<ext>` and
+# always ship an index; a shard-named file without one is a partial download.
+_SHARD_NAME_RE = re.compile(r"-\d+-of-\d+\.(safetensors|bin)$")
+
+
+def _indexed_shards_complete(folder: Path) -> bool | None:
+    """Check sharded weights against their index; ``None`` when no index exists."""
+    indexes = list(folder.rglob("*.index.json"))
+    if not indexes:
+        return None
+    for index in indexes:
+        try:
+            weight_map = json.loads(index.read_text()).get("weight_map") or {}
+        except (OSError, ValueError):
+            return False
+        shards = set(weight_map.values())
+        if not shards or any(not (index.parent / shard).is_file() for shard in shards):
+            return False
+    return True
+
+
+def _subdir_is_populated(base: str, subdir: str, needs_weights: bool) -> bool:
+    folder = Path(base) / subdir
+    if not folder.is_dir():
+        return False
+    if needs_weights:
+        indexed = _indexed_shards_complete(folder)
+        if indexed is not None:
+            return indexed
+        weights = [path for pattern in _WEIGHT_ARTIFACT_PATTERNS for path in folder.rglob(pattern)]
+        if not weights:
+            return False
+        return not any(_SHARD_NAME_RE.search(path.name) for path in weights)
+    return any((folder / name).is_file() for name in _TOKENIZER_ARTIFACT_NAMES)
+
+
+def _incomplete_stage_subdirs(
+    base: str,
+    subdirs: Sequence[str],
+    weight_subdirs: Collection[str] = (),
+) -> list[str]:
+    """Return the entries of ``subdirs`` without a POPULATED directory under ``base``.
+
+    Hub-snapshot paths use this stricter check: ``os.path.isdir`` alone accepts
+    a subfolder holding only config.json from an interrupted download, and the
+    warm-cache early return would then convert the Hub ID into a local path
+    vLLM cannot fetch missing weights for. ``weight_subdirs`` names the entries
+    that must contain a weight artifact, not merely any file.
+    """
+    return [
+        subdir for subdir in subdirs if not _subdir_is_populated(base, subdir, needs_weights=subdir in weight_subdirs)
+    ]
+
+
+def _resolve_model_to_local_path(
+    model: str,
+    required_subdirs: Sequence[str] = (),
+    weight_subdirs: Collection[str] = (),
+    *,
+    revision: str | None = None,
+    download_dir: str | None = None,
+) -> str:
+    """Resolve an HF Hub model ID to a local path that holds ``required_subdirs``.
+
+    ``snapshot_download(local_files_only=True)`` returns the snapshot root as
+    soon as *any* file of the repo is cached, even when the subfolders this
+    stage needs were never materialized. Joining a stage subdir onto such a
+    root produces a path that exists nowhere, and upstream ``EngineArgs``
+    forwards a non-directory ``model`` to HuggingFace as a repo id, which fails
+    with an ``HFValidationError`` about the cache path. Verify the subfolders
+    here and pull just the missing ones, so the join always lands on a real
+    directory or raises an error that names what is missing.
+
+    ``revision`` and ``download_dir`` mirror the engine args of the same name:
+    once the repo ID is replaced by a local path, downstream ModelConfig can no
+    longer correct either, so they must shape the snapshot selection here.
+    """
     if os.path.isdir(model):
         return model
 
+    # Keep the warm-cache path offline-friendly: no Hub round trip when the
+    # stage's subfolders are already there.
     try:
-        from huggingface_hub import snapshot_download
-
-        # Keep init path resolution offline-friendly.
-        return snapshot_download(model, local_files_only=True)
-    except Exception:
-        logger.warning(
-            "[stage_init] Could not resolve %s to local snapshot; using as-is",
-            model,
+        cached_root: str | None = hf_api().snapshot_download(
+            model, local_files_only=True, revision=revision, cache_dir=download_dir
         )
-        return model
+    except Exception:
+        cached_root = None
+    if cached_root is not None and not _incomplete_stage_subdirs(cached_root, required_subdirs, weight_subdirs):
+        return cached_root
+
+    # Cold cache, or a snapshot root whose stage subfolders were never (or
+    # only partially) downloaded: pull exactly the subfolders this stage asked
+    # for. snapshot_download resumes a partial subfolder for free.
+    allow_patterns = [f"{subdir.strip('/')}/*" for subdir in required_subdirs] or None
+    try:
+        resolved = hf_api().snapshot_download(
+            model, allow_patterns=allow_patterns, revision=revision, cache_dir=download_dir
+        )
+    except Exception as exc:
+        raise RuntimeError(
+            f"[stage_init] Could not resolve {model!r} to a local snapshot containing "
+            f"{sorted(required_subdirs)}: the download failed and "
+            + (
+                f"the cached snapshot {cached_root!r} is missing or incomplete for "
+                f"{sorted(_incomplete_stage_subdirs(cached_root, required_subdirs, weight_subdirs))}."
+                if cached_root is not None
+                else "nothing is cached locally."
+            )
+        ) from exc
+
+    missing = _incomplete_stage_subdirs(resolved, required_subdirs, weight_subdirs)
+    if missing:
+        raise RuntimeError(
+            f"[stage_init] Snapshot {resolved!r} for {model!r} has no populated {sorted(missing)} "
+            "subfolder; the stage cannot be initialized from it."
+        )
+    return resolved
 
 
 def _resolve_model_tokenizer_paths(model: str, engine_args: dict[str, Any]) -> str:
@@ -107,14 +234,66 @@ def _resolve_model_tokenizer_paths(model: str, engine_args: dict[str, Any]) -> s
     if model_subdir is None and tokenizer_subdir is None:
         return model
 
-    resolved_base = _resolve_model_to_local_path(model)
+    revision = engine_args.get("revision")
+    tokenizer_revision = engine_args.get("tokenizer_revision")
+    download_dir = engine_args.get("download_dir")
+    # A tokenizer pinned to a different revision cannot come from the model's
+    # snapshot; resolve it against its own. An empty subdir means the snapshot
+    # root and still needs its own revision.
+    split_tokenizer = tokenizer_subdir is not None and tokenizer_revision is not None and tokenizer_revision != revision
+
+    required_subdirs = [subdir for subdir in (model_subdir, tokenizer_subdir) if subdir]
+    model_required = [subdir for subdir in (model_subdir,) if subdir] if split_tokenizer else required_subdirs
+    weight_subdirs = frozenset(subdir for subdir in (model_subdir,) if subdir)
+    if is_runai_obj_uri(model):
+        # Object-storage URIs stay opaque until each stage builds its own
+        # ModelConfig, so the joins below are resolved by vLLM's streamer
+        # rather than by the local filesystem.
+        resolved_base = model
+        tokenizer_base = model
+    else:
+        resolved_base = _resolve_model_to_local_path(
+            model, model_required, weight_subdirs, revision=revision, download_dir=download_dir
+        )
+        # Reachable for a local model directory; the Hub branch above has
+        # already failed closed on a missing subfolder.
+        missing = _missing_stage_subdirs(resolved_base, model_required)
+        if missing:
+            raise RuntimeError(
+                f"[stage_init] Model directory {resolved_base!r} has no {sorted(missing)} "
+                "subfolder; the stage cannot be initialized from it."
+            )
+        if split_tokenizer:
+            # An empty subdir targets the snapshot root, which cannot be
+            # subset by allow_patterns; resolve the whole revision.
+            tokenizer_required = [tokenizer_subdir] if tokenizer_subdir else []
+            tokenizer_base = _resolve_model_to_local_path(
+                model, tokenizer_required, revision=tokenizer_revision, download_dir=download_dir
+            )
+            missing = _missing_stage_subdirs(tokenizer_base, tokenizer_required)
+            if missing:
+                raise RuntimeError(
+                    f"[stage_init] Tokenizer directory {tokenizer_base!r} has no {sorted(missing)} "
+                    "subfolder; the stage cannot be initialized from it."
+                )
+            if not tokenizer_required and not _subdir_is_populated(tokenizer_base, "", False):
+                # An empty subdir means the tokenizer lives at the snapshot
+                # root, so there is no subfolder for the check above to look
+                # at and any resolved root would otherwise pass. Require the
+                # vocabulary artifacts themselves.
+                raise RuntimeError(
+                    f"[stage_init] Tokenizer directory {tokenizer_base!r} holds none of "
+                    f"{list(_TOKENIZER_ARTIFACT_NAMES)}; the stage cannot be initialized from it."
+                )
+        else:
+            tokenizer_base = resolved_base
 
     if model_subdir:
         model = os.path.join(resolved_base, model_subdir)
         logger.info("[stage_init] Using model subdirectory: %s", model)
 
     if tokenizer_subdir is not None:
-        tokenizer_path = os.path.join(resolved_base, tokenizer_subdir) if tokenizer_subdir else resolved_base
+        tokenizer_path = os.path.join(tokenizer_base, tokenizer_subdir) if tokenizer_subdir else tokenizer_base
         engine_args["tokenizer"] = tokenizer_path
         logger.info("[stage_init] Using tokenizer from: %s", tokenizer_path)
     elif model_subdir and "tokenizer" not in engine_args:

vllm_omni/model_executor/models/minimax_music3/acoustic.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Stage-1 acoustic decoder for MiniMax Music 3.
 
 Consumes the AR stage's conditioning spans and returns 32 kHz stereo audio.
@@ -235,7 +235,11 @@ class MiniMaxMusic3AcousticForConditionalGeneration(nn.Module):
         would produce audio that sounds plausible but is wrong, which is far
         worse than a startup failure.
         """
-        root = resolve_repo_root(self.model_path)
+        root = resolve_repo_root(
+            self.model_path,
+            revision=self.vllm_config.model_config.revision,
+            download_dir=self.vllm_config.load_config.download_dir,
+        )
         encoder = self.dit.condition_encoder
         encoder.load_state_dict(load_component_state(root, CONDITION_ENCODER_DIR), strict=True)
         transformer_state = remap_transformer_state(load_component_state(root, TRANSFORMER_DIR))

vllm_omni/model_executor/models/minimax_music3/weights.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Component weight loading for MiniMax Music 3.
 
 The checkpoint is a multi-component repo. The Qwen3 backbone is a normal
@@ -35,26 +35,199 @@ CONDITION_ENCODER_DIR = "condition_encoder"
 TRANSFORMER_DIR = "transformer"
 VOCODER_DIR = "vocoder"
 
+# Folder names that identify a repository root. This is deliberately a
+# name-only probe and NOT the set a stage must be able to load: see
+# ``_REQUIRED_COMPONENTS``.
 _ROOT_MARKERS = (DEPTH_DECODER_DIR, TRANSFORMER_DIR, VOCODER_DIR)
 
+# Every component a stage actually reads. ``condition_encoder/`` is loaded by
+# the acoustic stage, so a snapshot missing it is incomplete even though the
+# root markers above are all present.
+_REQUIRED_COMPONENTS = (DEPTH_DECODER_DIR, CONDITION_ENCODER_DIR, TRANSFORMER_DIR, VOCODER_DIR)
 
-def resolve_repo_root(model_path: str | os.PathLike[str]) -> Path:
+# The subset a stage reading the repo ROOT needs on disk. The AR backbone under
+# ``language_model/`` is pulled by vLLM's own loader, and the repo's top-level
+# ``.pth`` files and ``qwen_7B/`` folder belong to the reference implementation,
+# so neither is downloaded here.
+_COMPONENT_PATTERNS = (
+    "*.json",
+    f"{DEPTH_DECODER_DIR}/*",
+    f"{CONDITION_ENCODER_DIR}/*",
+    f"{TRANSFORMER_DIR}/*",
+    f"{VOCODER_DIR}/*",
+)
+
+
+def _component_is_complete(folder: Path) -> bool:
+    """True when *folder* holds a weight set ``load_component_state`` can read.
+
+    A single shard satisfies a bare glob, so an interrupted multi-shard
+    download would otherwise pass as complete. Hub shards are named
+    ``<stem>-<index>-of-<total>.safetensors``, which states how many pieces
+    to expect, so that run can be checked for gaps.
+
+    Only that naming is provably incomplete. ``load_component_state`` loads
+    whatever the same glob returns, so any other shard-named file is accepted
+    here: judging it incomplete would force a needless re-download and fail
+    closed offline on weights that are present and readable.
+    """
+    stem = Path(_COMPONENT_WEIGHT_NAME).stem
+    if (folder / _COMPONENT_WEIGHT_NAME).is_file():
+        return True
+
+    shards = list(folder.glob(f"{stem}-*.safetensors"))
+    if not shards:
+        return False
+
+    totals: set[int] = set()
+    seen: set[int] = set()
+    for shard in shards:
+        marker = shard.name[len(stem) + 1 : -len(".safetensors")]
+        index_text, separator, total_text = marker.partition("-of-")
+        if not separator or not index_text.isdigit() or not total_text.isdigit():
+            continue
+        totals.add(int(total_text))
+        seen.add(int(index_text))
+    if not totals:
+        # No piece announces a total, so there is nothing to be short of.
+        return True
+    if len(totals) != 1:
+        return False
+    return seen == set(range(1, totals.pop() + 1))
+
+
+def _snapshot_has_components(root: Path) -> bool:
+    """True when every required component folder holds its complete weights."""
+    return all(_component_is_complete(root / component) for component in _REQUIRED_COMPONENTS)
+
+
+def _download_component_snapshot(
+    model: str,
+    *,
+    revision: str | None = None,
+    download_dir: str | None = None,
+) -> Path:
+    """Resolve an HF repo id to a snapshot holding the component folders.
+
+    Cache-first, so a warm cache needs no Hub round trip. Depending on the
+    huggingface_hub version, ``local_files_only=True`` either raises on a
+    partial cache or returns the partial snapshot root unchanged; both
+    variants must fall through to the online call, or a cold cache fails
+    startup without ever consulting the Hub.
+
+    ``revision`` and ``download_dir`` must be threaded through: resolving the
+    components against the default branch and the default cache while the AR
+    backbone was pinned elsewhere would silently mix weights from two commits.
+    """
+    from vllm.transformers_utils.repo_utils import hf_api
+
+    kwargs: dict[str, object] = {"allow_patterns": list(_COMPONENT_PATTERNS)}
+    if revision is not None:
+        kwargs["revision"] = revision
+    if download_dir is not None:
+        kwargs["cache_dir"] = download_dir
+
+    try:
+        cached = Path(hf_api().snapshot_download(model, local_files_only=True, **kwargs))
+    except Exception:
+        cached = None
+    if cached is not None and _snapshot_has_components(cached):
+        return cached
+    return Path(hf_api().snapshot_download(model, **kwargs))
+
+
+def _repo_id_and_revision_from_cache_path(path: Path) -> tuple[str | None, str | None]:
+    """Recover the Hub repo id and commit when *path* is inside an HF cache.
+
+    Stage init hands the talker the snapshot's ``language_model`` folder. On a
+    cold cache the sibling component folders were never downloaded, so the
+    marker walk fails even though the missing pieces are one snapshot_download
+    away. The cache layout is ``models--{org}--{name}/snapshots/{revision}/...``,
+    so the exact commit is recoverable too and must be reused: re-resolving the
+    repo id alone would fetch the default branch instead of the snapshot the
+    caller is already sitting in.
+    """
+    parts = path.parts
+    for position, part in enumerate(parts):
+        if not part.startswith("models--"):
+            continue
+        repo_id = part.removeprefix("models--").replace("--", "/")
+        revision = None
+        if position + 2 < len(parts) and parts[position + 1] == "snapshots":
+            revision = parts[position + 2]
+        return repo_id, revision
+    return None, None
+
+
+def resolve_repo_root(
+    model_path: str | os.PathLike[str],
+    *,
+    revision: str | None = None,
+    download_dir: str | None = None,
+) -> Path:
     """Find the multi-component repository root from a stage's model path.
 
     A stage that declares ``model_subdir`` sees the subfolder as its model
     path, so the root is usually the parent. Both are checked, and the search
     walks up a couple of levels for snapshot layouts.
 
+    The acoustic stage declares no ``model_subdir`` because it reads the root
+    itself, so on a hub-id deployment its model path is the repo id rather
+    than a directory. Resolve that to a local snapshot instead of treating it
+    as a relative path.
+
+    A cold cache is recoverable too: stage init pre-downloads only
+    ``language_model/`` and ``tokenizer/``, so the talker's model path exists
+    while the component folders do not. Fetch them here instead of failing.
+
+    A candidate is accepted only when its components are actually loadable.
+    Matching on folder names alone would accept a snapshot left weightless by
+    an interrupted download and defer the failure to ``load_component_state``,
+    past the point where the repair below could still run.
+
     Raises:
         FileNotFoundError: If no ancestor looks like the repository root.
     """
     path = Path(model_path).expanduser().resolve()
     candidates = [path, *path.parents[:3]]
-    for candidate in candidates:
-        if all((candidate / marker).is_dir() for marker in _ROOT_MARKERS):
+    marker_roots = [
+        candidate for candidate in candidates if all((candidate / marker).is_dir() for marker in _ROOT_MARKERS)
+    ]
+    for candidate in marker_roots:
+        if _snapshot_has_components(candidate):
             return candidate
+
+    # Either no ancestor carries the markers, or one does but is incomplete.
+    # Three recoverable cases: the path is a Hub repo id (the acoustic stage
+    # passes the id through), or it is a cache-snapshot subfolder whose sibling
+    # component folders were never downloaded, or a cached snapshot lost part
+    # of its weights. ``_download_component_snapshot`` is cache-first, so a
+    # warm cache stays offline either way.
+    cached_repo_id, cached_revision = _repo_id_and_revision_from_cache_path(path)
+    reference = str(model_path) if not path.exists() else cached_repo_id
+    if reference is not None:
+        try:
+            resolved = _download_component_snapshot(
+                reference,
+                revision=revision or cached_revision,
+                download_dir=download_dir,
+            )
+        except Exception as exc:
+            raise FileNotFoundError(
+                f"MiniMax Music 3 could not resolve {model_path!r} to a local repository root; "
+                f"expected folders {_REQUIRED_COMPONENTS}"
+            ) from exc
+        if _snapshot_has_components(resolved):
+            return resolved
+
+    if marker_roots:
+        raise FileNotFoundError(
+            f"MiniMax Music 3 found a repository root at {marker_roots[0]} but its components are incomplete; "
+            f"expected loadable weights under {_REQUIRED_COMPONENTS}"
+        )
     raise FileNotFoundError(
-        f"MiniMax Music 3 could not locate the repository root from {path}; expected sibling folders {_ROOT_MARKERS}"
+        f"MiniMax Music 3 could not locate the repository root from {path}; "
+        f"expected sibling folders {_REQUIRED_COMPONENTS}"
     )
 
 
@@ -99,7 +272,11 @@ def load_depth_decoder_weights(
     Raises:
         ValueError: If the audio embedding is absent or the wrong shape.
     """
-    root = resolve_repo_root(vllm_config.model_config.model)
+    root = resolve_repo_root(
+        vllm_config.model_config.model,
+        revision=vllm_config.model_config.revision,
+        download_dir=vllm_config.load_config.download_dir,
+    )
     state = load_component_state(root, DEPTH_DECODER_DIR)
 
     weight = state.pop(_AUDIO_EMBEDDING_KEY, None)