Skip to content

Commit

Permalink
[Core] Enable hf_transfer by default if available (#3817)
Browse files Browse the repository at this point in the history
  • Loading branch information
michaelfeil authored Apr 4, 2024
1 parent 294f8f6 commit 537ee25
Show file tree
Hide file tree
Showing 2 changed files with 42 additions and 0 deletions.
26 changes: 26 additions & 0 deletions tests/model_executor/weight_utils.py
Original file line number Diff line number Diff line change
@@ -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()
16 changes: 16 additions & 0 deletions vllm/model_executor/weight_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@
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
Expand All @@ -30,6 +31,21 @@
'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):
Expand Down

0 comments on commit 537ee25

Please sign in to comment.