[diffusion] fix: fix cache-first fast path accepting a metadata-only snapshot (#34121)
This commit is contained in:
@@ -77,6 +77,22 @@ def _model_hub_name() -> str:
|
|||||||
return "ModelScope" if envs.SGLANG_USE_MODELSCOPE.get() else "Hugging Face Hub"
|
return "ModelScope" if envs.SGLANG_USE_MODELSCOPE.get() else "Hugging Face Hub"
|
||||||
|
|
||||||
|
|
||||||
|
def _is_revisionless_snapshot_root(local_path: str) -> bool:
|
||||||
|
"""Detect a resolved "snapshot" that is really the ``snapshots/`` parent.
|
||||||
|
|
||||||
|
An empty ``refs/<revision>`` makes the offline resolver join ``""`` onto
|
||||||
|
``snapshots/`` and return the parent, which holds only revision subdirectories.
|
||||||
|
The ``models--*`` folder above is required too, since a ``local_dir`` may
|
||||||
|
legitimately be named ``snapshots``.
|
||||||
|
"""
|
||||||
|
head, tail = os.path.split(os.path.normpath(local_path))
|
||||||
|
return tail == "snapshots" and os.path.basename(head).split("--")[0] in (
|
||||||
|
"models",
|
||||||
|
"datasets",
|
||||||
|
"spaces",
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
def _snapshot_has_files(
|
def _snapshot_has_files(
|
||||||
local_path: str,
|
local_path: str,
|
||||||
allow_patterns: Optional[Union[list[str], str]],
|
allow_patterns: Optional[Union[list[str], str]],
|
||||||
@@ -169,6 +185,24 @@ def _get_missing_declared_weight_components(model_path: str) -> list[str]:
|
|||||||
return missing_files
|
return missing_files
|
||||||
|
|
||||||
|
|
||||||
|
def _is_metadata_only_pipeline_snapshot(model_path: str) -> bool:
|
||||||
|
"""Detect a snapshot holding only pipeline metadata, with no component weights.
|
||||||
|
|
||||||
|
``maybe_download_model_index`` probes a repo by fetching just ``model_index.json``;
|
||||||
|
that single-file fetch materializes a full cache entry, so a later
|
||||||
|
``local_files_only`` snapshot resolves it as a hit — offline there is no remote file
|
||||||
|
list to tell "cached" from "fully cached", so completeness must be read off disk.
|
||||||
|
|
||||||
|
Requires *all* declared components missing, not any: a partially populated snapshot
|
||||||
|
is legitimate (``allow_patterns``-filtered fetch), so only total absence is
|
||||||
|
unambiguously the probe stub. No declarations means no evidence to act on.
|
||||||
|
"""
|
||||||
|
declared = _get_declared_weight_component_dirs(model_path)
|
||||||
|
if not declared:
|
||||||
|
return False
|
||||||
|
return len(_get_missing_declared_weight_components(model_path)) == len(declared)
|
||||||
|
|
||||||
|
|
||||||
def _check_index_files_for_missing_shards(
|
def _check_index_files_for_missing_shards(
|
||||||
model_path: str,
|
model_path: str,
|
||||||
) -> tuple[bool, list[str], list[str]]:
|
) -> tuple[bool, list[str], list[str]]:
|
||||||
@@ -895,9 +929,31 @@ def maybe_download_model(
|
|||||||
local_files_only=True,
|
local_files_only=True,
|
||||||
max_workers=8,
|
max_workers=8,
|
||||||
)
|
)
|
||||||
|
if _is_revisionless_snapshot_root(local_path):
|
||||||
|
# A cache miss, so the download below re-resolves and rewrites the ref.
|
||||||
|
raise LocalEntryNotFoundError(
|
||||||
|
f"Cached ref for {model_name_or_path} is corrupt: resolved to the "
|
||||||
|
f"snapshots parent {local_path!r} instead of a revision directory."
|
||||||
|
)
|
||||||
if not force_diffusers_model:
|
if not force_diffusers_model:
|
||||||
return str(local_path)
|
# maybe_download_model_index's model_index.json fetch materializes a full
|
||||||
if is_lora or _verify_diffusers_model_complete(local_path):
|
# cache entry, so this resolve reports that stub as a hit; returning it
|
||||||
|
# would skip the download. LoRA repos declare no components.
|
||||||
|
if not is_lora and _is_metadata_only_pipeline_snapshot(local_path):
|
||||||
|
if not download:
|
||||||
|
raise ValueError(
|
||||||
|
f"Model {model_name_or_path} found in cache but only contains "
|
||||||
|
"pipeline metadata (no component weights) and download=False."
|
||||||
|
)
|
||||||
|
logger.info(
|
||||||
|
"Cached snapshot for %s only contains pipeline metadata, "
|
||||||
|
"will download component weights from %s",
|
||||||
|
model_name_or_path,
|
||||||
|
_model_hub_name(),
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
return str(local_path)
|
||||||
|
elif is_lora or _verify_diffusers_model_complete(local_path):
|
||||||
if not is_lora:
|
if not is_lora:
|
||||||
is_valid, cleanup_performed = _ci_validate_diffusers_model(local_path)
|
is_valid, cleanup_performed = _ci_validate_diffusers_model(local_path)
|
||||||
if not is_valid:
|
if not is_valid:
|
||||||
|
|||||||
@@ -7,7 +7,9 @@ from huggingface_hub.errors import LocalEntryNotFoundError
|
|||||||
from sglang.multimodal_gen.runtime.utils import hf_diffusers_utils
|
from sglang.multimodal_gen.runtime.utils import hf_diffusers_utils
|
||||||
from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import (
|
from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import (
|
||||||
_check_index_files_for_missing_shards,
|
_check_index_files_for_missing_shards,
|
||||||
|
_is_revisionless_snapshot_root,
|
||||||
_verify_diffusers_model_complete,
|
_verify_diffusers_model_complete,
|
||||||
|
maybe_download_model,
|
||||||
)
|
)
|
||||||
from sglang.srt.environ import envs
|
from sglang.srt.environ import envs
|
||||||
|
|
||||||
@@ -78,6 +80,179 @@ def test_diffusers_cache_validation_checks_declared_component_shards(tmp_path):
|
|||||||
assert "transformer" in checked_subdirs
|
assert "transformer" in checked_subdirs
|
||||||
|
|
||||||
|
|
||||||
|
def _populate_components(root, components, *, weights=True):
|
||||||
|
for name in components:
|
||||||
|
(root / name).mkdir(parents=True, exist_ok=True)
|
||||||
|
if weights:
|
||||||
|
(root / name / "model.safetensors").write_bytes(b"weights")
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def recording_snapshot_download(monkeypatch):
|
||||||
|
"""Stub ``snapshot_download`` to return ``path``, recording each call as an
|
||||||
|
offline ``"probe"`` or a real ``"download"``."""
|
||||||
|
calls = []
|
||||||
|
|
||||||
|
def factory(path):
|
||||||
|
def fake_snapshot_download(**kwargs):
|
||||||
|
calls.append("probe" if kwargs.get("local_files_only") else "download")
|
||||||
|
return str(path)
|
||||||
|
|
||||||
|
monkeypatch.setattr(
|
||||||
|
hf_diffusers_utils, "snapshot_download", fake_snapshot_download
|
||||||
|
)
|
||||||
|
return calls
|
||||||
|
|
||||||
|
return factory
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.parametrize(
|
||||||
|
"relative_path, expected",
|
||||||
|
[
|
||||||
|
("models--org--repo/snapshots", True),
|
||||||
|
("models--org--repo/snapshots/", True),
|
||||||
|
("datasets--org--repo/snapshots", True),
|
||||||
|
# A healthy entry resolves to snapshots/<commit_sha>.
|
||||||
|
("models--org--repo/snapshots/" + "a" * 40, False),
|
||||||
|
# A local_dir may legitimately be named "snapshots".
|
||||||
|
("mymodels/snapshots", False),
|
||||||
|
("mymodels/snapshots/inner", False),
|
||||||
|
("models--org--repo", False),
|
||||||
|
],
|
||||||
|
)
|
||||||
|
def test_revisionless_snapshot_root_detection(tmp_path, relative_path, expected):
|
||||||
|
assert _is_revisionless_snapshot_root(str(tmp_path / relative_path)) is expected
|
||||||
|
|
||||||
|
|
||||||
|
def test_corrupt_ref_resolving_to_snapshots_parent_is_a_cache_miss(
|
||||||
|
recording_snapshot_download, tmp_path
|
||||||
|
):
|
||||||
|
"""The snapshots/ parent holds only revision subdirs, so it is not a hit."""
|
||||||
|
snapshots_parent = tmp_path / "models--org--repo" / "snapshots"
|
||||||
|
(snapshots_parent / ("a" * 40)).mkdir(parents=True)
|
||||||
|
calls = recording_snapshot_download(snapshots_parent)
|
||||||
|
|
||||||
|
with pytest.raises(ValueError, match="not found in local cache"):
|
||||||
|
maybe_download_model("org/repo", download=False)
|
||||||
|
|
||||||
|
assert calls == ["probe"]
|
||||||
|
|
||||||
|
|
||||||
|
def test_local_dir_named_snapshots_is_not_treated_as_corrupt(
|
||||||
|
recording_snapshot_download, tmp_path
|
||||||
|
):
|
||||||
|
"""The guard requires a models--*/ folder above, so this stays a valid hit."""
|
||||||
|
local_dir = tmp_path / "mymodels" / "snapshots"
|
||||||
|
local_dir.mkdir(parents=True)
|
||||||
|
_write_model_index(local_dir)
|
||||||
|
_populate_components(local_dir, ("text_encoder", "transformer", "vae"))
|
||||||
|
calls = recording_snapshot_download(local_dir)
|
||||||
|
|
||||||
|
result = maybe_download_model("org/repo", local_dir=str(local_dir), download=False)
|
||||||
|
|
||||||
|
assert result == str(local_dir)
|
||||||
|
assert calls == ["probe"]
|
||||||
|
|
||||||
|
|
||||||
|
def test_local_path_is_returned_as_given(tmp_path):
|
||||||
|
"""A local path is not a repo id, so it is returned without any download."""
|
||||||
|
_write_model_index(tmp_path)
|
||||||
|
_populate_components(tmp_path, ("text_encoder", "transformer", "vae"))
|
||||||
|
|
||||||
|
assert maybe_download_model(str(tmp_path), download=False) == str(tmp_path)
|
||||||
|
|
||||||
|
|
||||||
|
def test_metadata_only_local_path_is_still_returned_as_given(tmp_path):
|
||||||
|
"""Falling through would download the filesystem path as a repo id and fail."""
|
||||||
|
_write_model_index(tmp_path)
|
||||||
|
|
||||||
|
assert maybe_download_model(str(tmp_path), download=False) == str(tmp_path)
|
||||||
|
|
||||||
|
|
||||||
|
def test_metadata_only_cached_snapshot_is_not_a_usable_hit(
|
||||||
|
recording_snapshot_download, tmp_path
|
||||||
|
):
|
||||||
|
"""The probe stub resolves as an offline hit but has no component weights."""
|
||||||
|
_write_model_index(tmp_path)
|
||||||
|
calls = recording_snapshot_download(tmp_path)
|
||||||
|
|
||||||
|
with pytest.raises(ValueError, match="only contains pipeline metadata"):
|
||||||
|
maybe_download_model("org/repo", download=False)
|
||||||
|
|
||||||
|
assert calls == ["probe"]
|
||||||
|
|
||||||
|
|
||||||
|
def test_metadata_only_cached_snapshot_falls_through_to_one_download(
|
||||||
|
recording_snapshot_download, tmp_path
|
||||||
|
):
|
||||||
|
"""Exactly one download -- no duplicate fetch, and no force_download retry."""
|
||||||
|
_write_model_index(tmp_path)
|
||||||
|
calls = recording_snapshot_download(tmp_path)
|
||||||
|
|
||||||
|
assert maybe_download_model("org/repo") == str(tmp_path)
|
||||||
|
assert calls == ["probe", "download"]
|
||||||
|
|
||||||
|
|
||||||
|
def test_complete_cached_snapshot_is_served_without_download(
|
||||||
|
recording_snapshot_download, tmp_path
|
||||||
|
):
|
||||||
|
_write_model_index(tmp_path)
|
||||||
|
_populate_components(tmp_path, ("text_encoder", "transformer", "vae"))
|
||||||
|
calls = recording_snapshot_download(tmp_path)
|
||||||
|
|
||||||
|
assert maybe_download_model("org/repo") == str(tmp_path)
|
||||||
|
assert calls == ["probe"]
|
||||||
|
|
||||||
|
|
||||||
|
def test_partially_populated_cached_snapshot_is_served_without_download(
|
||||||
|
recording_snapshot_download, tmp_path
|
||||||
|
):
|
||||||
|
"""An allow_patterns fetch leaves components absent, so only a total absence of
|
||||||
|
weights counts as the stub."""
|
||||||
|
_write_model_index(tmp_path)
|
||||||
|
_populate_components(tmp_path, ("vae",))
|
||||||
|
calls = recording_snapshot_download(tmp_path)
|
||||||
|
|
||||||
|
assert maybe_download_model("org/repo") == str(tmp_path)
|
||||||
|
assert calls == ["probe"]
|
||||||
|
|
||||||
|
|
||||||
|
def test_cached_component_repo_without_model_index_is_served_without_download(
|
||||||
|
recording_snapshot_download, tmp_path
|
||||||
|
):
|
||||||
|
"""Declares nothing, so the stub check must not match on 0 missing == 0."""
|
||||||
|
(tmp_path / "model.safetensors").write_bytes(b"weights")
|
||||||
|
calls = recording_snapshot_download(tmp_path)
|
||||||
|
|
||||||
|
assert maybe_download_model("org/repo") == str(tmp_path)
|
||||||
|
assert calls == ["probe"]
|
||||||
|
|
||||||
|
|
||||||
|
def test_metadata_only_cached_lora_snapshot_is_a_usable_hit(
|
||||||
|
recording_snapshot_download, tmp_path
|
||||||
|
):
|
||||||
|
"""LoRA repos declare no components, so they can never be the stub."""
|
||||||
|
_write_model_index(tmp_path)
|
||||||
|
calls = recording_snapshot_download(tmp_path)
|
||||||
|
|
||||||
|
assert maybe_download_model("org/repo", is_lora=True) == str(tmp_path)
|
||||||
|
assert calls == ["probe"]
|
||||||
|
|
||||||
|
|
||||||
|
def test_force_diffusers_model_stub_keeps_its_existing_path(
|
||||||
|
recording_snapshot_download, tmp_path
|
||||||
|
):
|
||||||
|
"""Already rejected by _verify_diffusers_model_complete, so its path is
|
||||||
|
unchanged: download, then the force_download retry."""
|
||||||
|
_write_model_index(tmp_path)
|
||||||
|
calls = recording_snapshot_download(tmp_path)
|
||||||
|
|
||||||
|
with pytest.raises(ValueError, match="still incomplete after forced re-download"):
|
||||||
|
maybe_download_model("org/repo", force_diffusers_model=True)
|
||||||
|
|
||||||
|
assert calls == ["probe", "download", "download"]
|
||||||
|
|
||||||
|
|
||||||
def test_modelscope_file_download_preserves_local_dir(monkeypatch, tmp_path):
|
def test_modelscope_file_download_preserves_local_dir(monkeypatch, tmp_path):
|
||||||
calls = []
|
calls = []
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user