feat: add HuggingFace mirror support for VM providers (#278)
Add support for HuggingFace mirror (hf-mirror.com) to improve download speeds in regions where huggingface.co access is slow. - Support HF_ENDPOINT environment variable detection - Automatically switch to hf-mirror.com when HF_ENDPOINT is set - Apply to Docker, VMware, and VirtualBox providers - Maintain backward compatibility with default huggingface.co URLs Users can now set HF_ENDPOINT=https://hf-mirror.com to use the mirror.
This commit is contained in:
@@ -34,6 +34,12 @@ def _download_vm(vms_dir: str):
|
|||||||
logger.info("Downloading the virtual machine image...")
|
logger.info("Downloading the virtual machine image...")
|
||||||
downloaded_size = 0
|
downloaded_size = 0
|
||||||
|
|
||||||
|
# Check for HF_ENDPOINT environment variable and replace domain if set to hf-mirror.com
|
||||||
|
hf_endpoint = os.environ.get('HF_ENDPOINT')
|
||||||
|
if hf_endpoint and 'hf-mirror.com' in hf_endpoint:
|
||||||
|
URL = URL.replace('huggingface.co', 'hf-mirror.com')
|
||||||
|
logger.info(f"Using HF mirror: {URL}")
|
||||||
|
|
||||||
downloaded_file_name = DOWNLOADED_FILE_NAME
|
downloaded_file_name = DOWNLOADED_FILE_NAME
|
||||||
|
|
||||||
os.makedirs(vms_dir, exist_ok=True)
|
os.makedirs(vms_dir, exist_ok=True)
|
||||||
@@ -113,6 +119,13 @@ class DockerVMManager(VMManager):
|
|||||||
URL = UBUNTU_X86_URL
|
URL = UBUNTU_X86_URL
|
||||||
elif os_type == "Windows":
|
elif os_type == "Windows":
|
||||||
URL = WINDOWS_X86_URL
|
URL = WINDOWS_X86_URL
|
||||||
|
|
||||||
|
# Check for HF_ENDPOINT environment variable and replace domain if set to hf-mirror.com
|
||||||
|
hf_endpoint = os.environ.get('HF_ENDPOINT')
|
||||||
|
if hf_endpoint and 'hf-mirror.com' in hf_endpoint:
|
||||||
|
URL = URL.replace('huggingface.co', 'hf-mirror.com')
|
||||||
|
logger.info(f"Using HF mirror: {URL}")
|
||||||
|
|
||||||
DOWNLOADED_FILE_NAME = URL.split('/')[-1]
|
DOWNLOADED_FILE_NAME = URL.split('/')[-1]
|
||||||
|
|
||||||
if DOWNLOADED_FILE_NAME.endswith(".zip"):
|
if DOWNLOADED_FILE_NAME.endswith(".zip"):
|
||||||
|
|||||||
@@ -57,6 +57,12 @@ def _install_vm(vm_name, vms_dir, downloaded_file_name, original_vm_name="Ubuntu
|
|||||||
else:
|
else:
|
||||||
raise Exception("Unsupported platform or architecture.")
|
raise Exception("Unsupported platform or architecture.")
|
||||||
|
|
||||||
|
# Check for HF_ENDPOINT environment variable and replace domain if set to hf-mirror.com
|
||||||
|
hf_endpoint = os.environ.get('HF_ENDPOINT')
|
||||||
|
if hf_endpoint and 'hf-mirror.com' in hf_endpoint:
|
||||||
|
url = url.replace('huggingface.co', 'hf-mirror.com')
|
||||||
|
logger.info(f"Using HF mirror: {url}")
|
||||||
|
|
||||||
# Download the virtual machine image
|
# Download the virtual machine image
|
||||||
logger.info("Downloading the virtual machine image...")
|
logger.info("Downloading the virtual machine image...")
|
||||||
downloaded_size = 0
|
downloaded_size = 0
|
||||||
|
|||||||
@@ -134,6 +134,13 @@ def _install_vm(vm_name, vms_dir, downloaded_file_name, os_type, original_vm_nam
|
|||||||
elif os_type == "Windows":
|
elif os_type == "Windows":
|
||||||
if platform.machine().lower() in ['amd64', 'x86_64']:
|
if platform.machine().lower() in ['amd64', 'x86_64']:
|
||||||
URL = WINDOWS_X86_URL
|
URL = WINDOWS_X86_URL
|
||||||
|
|
||||||
|
# Check for HF_ENDPOINT environment variable and replace domain if set to hf-mirror.com
|
||||||
|
hf_endpoint = os.environ.get('HF_ENDPOINT')
|
||||||
|
if hf_endpoint and 'hf-mirror.com' in hf_endpoint:
|
||||||
|
URL = URL.replace('huggingface.co', 'hf-mirror.com')
|
||||||
|
logger.info(f"Using HF mirror: {URL}")
|
||||||
|
|
||||||
DOWNLOADED_FILE_NAME = URL.split('/')[-1]
|
DOWNLOADED_FILE_NAME = URL.split('/')[-1]
|
||||||
downloaded_file_name = DOWNLOADED_FILE_NAME
|
downloaded_file_name = DOWNLOADED_FILE_NAME
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user