mirror of
https://git.datalinker.icu/vllm-project/vllm.git
synced 2025-12-09 22:25:32 +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
|
||||
|
||||
import filelock
|
||||
import huggingface_hub.constants
|
||||
import numpy as np
|
||||
import torch
|
||||
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/"
|
||||
|
||||
|
||||
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):
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user