[Core] Enable hf_transfer by default if available (#3817)

This commit is contained in:
Michael Feil 2024-04-03 21:02:43 -07:00 committed by GitHub
parent 294f8f6665
commit 537ee25f43
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 42 additions and 0 deletions

View 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()

View File

@ -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):