diff --git a/tests/engine/test_async_omni_engine_stage_init.py b/tests/engine/test_async_omni_engine_stage_init.py index 23e56d95b..c2bbc2ed6 100644 --- a/tests/engine/test_async_omni_engine_stage_init.py +++ b/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 diff --git a/tests/engine/test_stage_engine_args.py b/tests/engine/test_stage_engine_args.py index 6afde7469..0a54e2e80 100644 --- a/tests/engine/test_stage_engine_args.py +++ b/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 diff --git a/tests/model_executor/models/test_minimax_music3_repo_root.py b/tests/model_executor/models/test_minimax_music3_repo_root.py new file mode 100644 index 000000000..182635ab7 --- /dev/null +++ b/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/`` 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 ``-*.safetensors`` and reads whatever it finds, so a + component whose weights do not use the ``--of-`` 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 == [] diff --git a/tools/pre_commit/check_forbidden_imports.py b/tools/pre_commit/check_forbidden_imports.py index a3f5e320a..926e102fa 100644 --- a/tools/pre_commit/check_forbidden_imports.py +++ b/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", diff --git a/vllm_omni/engine/stage_init_utils.py b/vllm_omni/engine/stage_init_utils.py index 6895bdad8..b1589aef5 100644 --- a/vllm_omni/engine/stage_init_utils.py +++ b/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 `-NNNNN-of-NNNNN.` 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: diff --git a/vllm_omni/model_executor/models/minimax_music3/acoustic.py b/vllm_omni/model_executor/models/minimax_music3/acoustic.py index 879185a57..92fa332d0 100644 --- a/vllm_omni/model_executor/models/minimax_music3/acoustic.py +++ b/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)) diff --git a/vllm_omni/model_executor/models/minimax_music3/weights.py b/vllm_omni/model_executor/models/minimax_music3/weights.py index 2ab643c65..b019ff461 100644 --- a/vllm_omni/model_executor/models/minimax_music3/weights.py +++ b/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 + ``--of-.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)