Fix URL when uploading to proxy (#2167)

* Fix URL when uploading to proxy

* fix multi-part LFS upload

Author: Wauplin

src/huggingface_hub/_commit_api.py

@@ -399,7 +399,7 @@ def _upload_lfs_files(
     def _wrapped_lfs_upload(batch_action) -> None:
         try:
             operation = oid2addop[batch_action["oid"]]
-            lfs_upload(operation=operation, lfs_batch_action=batch_action, headers=headers)
+            lfs_upload(operation=operation, lfs_batch_action=batch_action, headers=headers, endpoint=endpoint)
         except Exception as exc:
             raise RuntimeError(f"Error while uploading '{operation.path_in_repo}' to the Hub.") from exc
 

src/huggingface_hub/hf_api.py

@@ -82,8 +82,6 @@
     deserialize_event,
 )
 from .constants import (
-    _HF_DEFAULT_ENDPOINT,
-    _HF_DEFAULT_STAGING_ENDPOINT,
     DEFAULT_ETAG_TIMEOUT,
     DEFAULT_REQUEST_TIMEOUT,
     DEFAULT_REVISION,
@@ -123,6 +121,7 @@
     build_hf_headers,
     experimental,
     filter_repo_objects,
+    fix_hf_endpoint_in_url,
     get_session,
     hf_raise_for_status,
     logging,
@@ -431,14 +430,7 @@ class RepoUrl(str):
     """
 
     def __new__(cls, url: Any, endpoint: Optional[str] = None):
-        # check if a proxy has been set => if yes, update the returned URL to use the proxy
-        if ENDPOINT not in (_HF_DEFAULT_ENDPOINT, _HF_DEFAULT_STAGING_ENDPOINT):
-            url = url.replace(_HF_DEFAULT_ENDPOINT, ENDPOINT)
-            url = url.replace(_HF_DEFAULT_STAGING_ENDPOINT, ENDPOINT)
-        if endpoint not in (None, _HF_DEFAULT_ENDPOINT, _HF_DEFAULT_STAGING_ENDPOINT):
-            url = url.replace(_HF_DEFAULT_ENDPOINT, endpoint)
-            url = url.replace(_HF_DEFAULT_STAGING_ENDPOINT, endpoint)
-
+        url = fix_hf_endpoint_in_url(url, endpoint=endpoint)
         return super(RepoUrl, cls).__new__(cls, url)
 
     def __init__(self, url: Any, endpoint: Optional[str] = None) -> None:

src/huggingface_hub/lfs.py

@@ -31,6 +31,7 @@
 
 from .utils import (
     build_hf_headers,
+    fix_hf_endpoint_in_url,
     get_session,
     hf_raise_for_status,
     http_backoff,
@@ -193,6 +194,7 @@ def lfs_upload(
     lfs_batch_action: Dict,
     token: Optional[str] = None,
     headers: Optional[Dict[str, str]] = None,
+    endpoint: Optional[str] = None,
 ) -> None:
     """
     Handles uploading a given object to the Hub with the LFS protocol.
@@ -230,22 +232,24 @@ def lfs_upload(
     # 2. Upload file (either single part or multi-part)
     header = upload_action.get("header", {})
     chunk_size = header.get("chunk_size")
+    upload_url = fix_hf_endpoint_in_url(upload_action["href"], endpoint=endpoint)
     if chunk_size is not None:
         try:
             chunk_size = int(chunk_size)
         except (ValueError, TypeError):
             raise ValueError(
                 f"Malformed response from LFS batch endpoint: `chunk_size` should be an integer. Got '{chunk_size}'."
             )
-        _upload_multi_part(operation=operation, header=header, chunk_size=chunk_size, upload_url=upload_action["href"])
+        _upload_multi_part(operation=operation, header=header, chunk_size=chunk_size, upload_url=upload_url)
     else:
-        _upload_single_part(operation=operation, upload_url=upload_action["href"])
+        _upload_single_part(operation=operation, upload_url=upload_url)
 
     # 3. Verify upload went well
     if verify_action is not None:
         _validate_lfs_action(verify_action)
+        verify_url = fix_hf_endpoint_in_url(verify_action["href"], endpoint)
         verify_resp = get_session().post(
-            verify_action["href"],
+            verify_url,
             headers=build_hf_headers(token=token, headers=headers),
             json={"oid": operation.upload_info.sha256.hex(), "size": operation.upload_info.size},
         )

src/huggingface_hub/utils/__init__.py

@@ -47,7 +47,14 @@
 from ._git_credential import list_credential_helpers, set_git_credential, unset_git_credential
 from ._headers import LocalTokenNotFoundError, build_hf_headers, get_token_to_send
 from ._hf_folder import HfFolder
-from ._http import OfflineModeIsEnabled, configure_http_backend, get_session, http_backoff, reset_sessions
+from ._http import (
+    OfflineModeIsEnabled,
+    configure_http_backend,
+    fix_hf_endpoint_in_url,
+    get_session,
+    http_backoff,
+    reset_sessions,
+)
 from ._pagination import paginate
 from ._paths import IGNORE_GIT_FOLDER_PATTERNS, filter_repo_objects
 from ._runtime import (

src/huggingface_hub/utils/_http.py

@@ -21,7 +21,7 @@
 import uuid
 from functools import lru_cache
 from http import HTTPStatus
-from typing import Callable, Tuple, Type, Union
+from typing import Callable, Optional, Tuple, Type, Union
 
 import requests
 from requests import Response
@@ -306,3 +306,16 @@ def http_backoff(
 
         # Update sleep time for next retry
         sleep_time = min(max_wait_time, sleep_time * 2)  # Exponential backoff
+
+
+def fix_hf_endpoint_in_url(url: str, endpoint: Optional[str]) -> str:
+    """Replace the default endpoint in a URL by a custom one.
+
+    This is useful when using a proxy and the Hugging Face Hub returns a URL with the default endpoint.
+    """
+    endpoint = endpoint or constants.ENDPOINT
+    # check if a proxy has been set => if yes, update the returned URL to use the proxy
+    if endpoint not in (None, constants._HF_DEFAULT_ENDPOINT, constants._HF_DEFAULT_STAGING_ENDPOINT):
+        url = url.replace(constants._HF_DEFAULT_ENDPOINT, endpoint)
+        url = url.replace(constants._HF_DEFAULT_STAGING_ENDPOINT, endpoint)
+    return url

tests/test_utils_http.py

@@ -3,17 +3,19 @@
 import time
 import unittest
 from multiprocessing import Process, Queue
-from typing import Generator
+from typing import Generator, Optional
 from unittest.mock import Mock, call, patch
 from uuid import UUID
 
+import pytest
 import requests
 from requests import ConnectTimeout, HTTPError
 
 from huggingface_hub.constants import ENDPOINT
 from huggingface_hub.utils._http import (
     OfflineModeIsEnabled,
     configure_http_backend,
+    fix_hf_endpoint_in_url,
     get_session,
     http_backoff,
     reset_sessions,
@@ -292,3 +294,20 @@ def _is_uuid(string: str) -> bool:
     except ValueError:
         return False
     return str(uuid_obj) == string
+
+
+@pytest.mark.parametrize(
+    ("base_url", "endpoint", "expected_url"),
+    [
+        # Staging url => unchanged
+        ("https://hub-ci.huggingface.co/resolve/...", None, "https://hub-ci.huggingface.co/resolve/..."),
+        # Prod url => unchanged
+        ("https://huggingface.co/resolve/...", None, "https://huggingface.co/resolve/..."),
+        # Custom endpoint + staging url => fixed
+        ("https://hub-ci.huggingface.co/api/models", "https://mirror.co", "https://mirror.co/api/models"),
+        # Custom endpoint + prod url => fixed
+        ("https://huggingface.co/api/models", "https://mirror.co", "https://mirror.co/api/models"),
+    ],
+)
+def test_fix_hf_endpoint_in_url(base_url: str, endpoint: Optional[str], expected_url: str) -> None:
+    assert fix_hf_endpoint_in_url(base_url, endpoint) == expected_url