weight_utils.py 2.09 KB
Newer Older
1
2
# SPDX-License-Identifier: Apache-2.0

3
import os
4
import tempfile
5
6
7

import huggingface_hub.constants
import pytest
8
from huggingface_hub.utils import LocalEntryNotFoundError
9

10
11
from vllm.model_executor.model_loader.weight_utils import (
    download_weights_from_hf, enable_hf_transfer)
12
from ..utils import models_path_prefix
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30


def test_hf_transfer_auto_activation():
    if "HF_HUB_ENABLE_HF_TRANSFER" in os.environ:
        # in case it is already set, we can't test the auto activation
        pytest.skip(
            "HF_HUB_ENABLE_HF_TRANSFER is set, can't test auto activation")
    enable_hf_transfer()
    try:
        # enable hf hub transfer if available
        import hf_transfer  # type: ignore # noqa
        HF_TRANFER_ACTIVE = True
    except ImportError:
        HF_TRANFER_ACTIVE = False
    assert (huggingface_hub.constants.HF_HUB_ENABLE_HF_TRANSFER ==
            HF_TRANFER_ACTIVE)


31
32
33
34
35
36
def test_download_weights_from_hf():
    with tempfile.TemporaryDirectory() as tmpdir:
        # assert LocalEntryNotFoundError error is thrown
        # if offline is set and model is not cached
        huggingface_hub.constants.HF_HUB_OFFLINE = True
        with pytest.raises(LocalEntryNotFoundError):
37
            download_weights_from_hf(os.path.join(models_path_prefix, "facebook/opt-125m"),
38
39
40
41
42
                                     allow_patterns=["*.safetensors", "*.bin"],
                                     cache_dir=tmpdir)

        # download the model
        huggingface_hub.constants.HF_HUB_OFFLINE = False
43
        download_weights_from_hf(os.path.join(models_path_prefix, "facebook/opt-125m"),
44
45
46
47
48
49
                                 allow_patterns=["*.safetensors", "*.bin"],
                                 cache_dir=tmpdir)

        # now it should work offline
        huggingface_hub.constants.HF_HUB_OFFLINE = True
        assert download_weights_from_hf(
50
            os.path.join(models_path_prefix, "facebook/opt-125m"),
51
52
53
54
            allow_patterns=["*.safetensors", "*.bin"],
            cache_dir=tmpdir) is not None


55
56
if __name__ == "__main__":
    test_hf_transfer_auto_activation()
57
    test_download_weights_from_hf()