From 5b66a14ffca7209abfda010e7970ec1f6a146964 Mon Sep 17 00:00:00 2001
From: Michael Feil <63565275+michaelfeil@users.noreply.github.com>
Date: Wed, 3 Apr 2024 21:02:43 -0700
Subject: [PATCH] [Core] Enable hf_transfer by default if available (#3817)

---
 tests/model_executor/weight_utils.py | 26 ++++++++++++++++++++++++++
 vllm/model_executor/weight_utils.py  | 16 ++++++++++++++++
 2 files changed, 42 insertions(+)
 create mode 100644 tests/model_executor/weight_utils.py

diff --git a/tests/model_executor/weight_utils.py b/tests/model_executor/weight_utils.py
new file mode 100644
index 0000000000000..3154f2826d10c
--- /dev/null
+++ b/tests/model_executor/weight_utils.py
@@ -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()
diff --git a/vllm/model_executor/weight_utils.py b/vllm/model_executor/weight_utils.py
index 2c0dd8ce76036..0961478930d74 100644
--- a/vllm/model_executor/weight_utils.py
+++ b/vllm/model_executor/weight_utils.py
@@ -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
@@ -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):