mirror of
https://git.datalinker.icu/vllm-project/vllm.git
synced 2025-12-10 06:45:01 +08:00
[Core] Enable hf_transfer by default if available (#3817)
This commit is contained in:
parent
294f8f6665
commit
537ee25f43
26
tests/model_executor/weight_utils.py
Normal file
26
tests/model_executor/weight_utils.py
Normal file
@ -0,0 +1,26 @@
|
|||||||
|
import os
|
||||||
|
|
||||||
|
import huggingface_hub.constants
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
from vllm.model_executor.weight_utils import enable_hf_transfer
|
||||||
|
|
||||||
|
|
||||||
|
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)
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
test_hf_transfer_auto_activation()
|
||||||
@ -8,6 +8,7 @@ from collections import defaultdict
|
|||||||
from typing import Any, Iterable, Iterator, List, Optional, Tuple
|
from typing import Any, Iterable, Iterator, List, Optional, Tuple
|
||||||
|
|
||||||
import filelock
|
import filelock
|
||||||
|
import huggingface_hub.constants
|
||||||
import numpy as np
|
import numpy as np
|
||||||
import torch
|
import torch
|
||||||
from huggingface_hub import HfFileSystem, snapshot_download
|
from huggingface_hub import HfFileSystem, snapshot_download
|
||||||
@ -30,6 +31,21 @@ temp_dir = os.environ.get('TMPDIR') or os.environ.get(
|
|||||||
'TEMP') or os.environ.get('TMP') or "/tmp/"
|
'TEMP') or os.environ.get('TMP') or "/tmp/"
|
||||||
|
|
||||||
|
|
||||||
|
def enable_hf_transfer():
|
||||||
|
"""automatically activates hf_transfer
|
||||||
|
"""
|
||||||
|
if "HF_HUB_ENABLE_HF_TRANSFER" not in os.environ:
|
||||||
|
try:
|
||||||
|
# enable hf hub transfer if available
|
||||||
|
import hf_transfer # type: ignore # noqa
|
||||||
|
huggingface_hub.constants.HF_HUB_ENABLE_HF_TRANSFER = True
|
||||||
|
except ImportError:
|
||||||
|
pass
|
||||||
|
|
||||||
|
|
||||||
|
enable_hf_transfer()
|
||||||
|
|
||||||
|
|
||||||
class Disabledtqdm(tqdm):
|
class Disabledtqdm(tqdm):
|
||||||
|
|
||||||
def __init__(self, *args, **kwargs):
|
def __init__(self, *args, **kwargs):
|
||||||
|
|||||||
Loading…
x
Reference in New Issue
Block a user