Unverified Commit 537ee25f authored by Michael Feil's avatar Michael Feil Committed by GitHub
Browse files

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

parent 294f8f66
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 ...@@ -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( ...@@ -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):
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment