Avoid redundant call to the Xet connection info URL (#3534)

* Avoid redundant call to the Xet connection info URL

* authorization casing

* remove log

* thread safety

Author: Wauplin

src/huggingface_hub/utils/_xet.py

@@ -1,3 +1,4 @@
+import time
 from dataclasses import dataclass
 from enum import Enum
 from typing import Optional
@@ -8,6 +9,11 @@
 from . import get_session, hf_raise_for_status, validate_hf_hub_args
 
 
+XET_CONNECTION_INFO_SAFETY_PERIOD = 60  # seconds
+XET_CONNECTION_INFO_CACHE_SIZE = 1_000
+XET_CONNECTION_INFO_CACHE: dict[str, "XetConnectionInfo"] = {}
+
+
 class XetTokenType(str, Enum):
     READ = "read"
     WRITE = "write"
@@ -167,6 +173,9 @@ def _fetch_xet_connection_info_with_url(
     """
     Requests the xet connection info from the supplied URL. This includes the
     access token, expiration time, and endpoint to use for the xet storage service.
+
+    Result is cached to avoid redundant requests.
+
     Args:
         url: (`str`):
             The access token endpoint URL.
@@ -183,10 +192,44 @@ def _fetch_xet_connection_info_with_url(
         [`ValueError`](https://docs.python.org/3/library/exceptions.html#ValueError)
             If the Hub API response is improperly formatted.
     """
+    # Check cache first
+    cache_key = _cache_key(url, headers, params)
+    cached_info = XET_CONNECTION_INFO_CACHE.get(cache_key)
+    if cached_info is not None:
+        if not _is_expired(cached_info):
+            return cached_info
+
+    # Fetch from server
     resp = get_session().get(headers=headers, url=url, params=params)
     hf_raise_for_status(resp)
 
     metadata = parse_xet_connection_info_from_headers(resp.headers)  # type: ignore
     if metadata is None:
         raise ValueError("Xet headers have not been correctly set by the server.")
+
+    # Delete expired cache entries
+    for k, v in list(XET_CONNECTION_INFO_CACHE.items()):
+        if _is_expired(v):
+            XET_CONNECTION_INFO_CACHE.pop(k, None)
+
+    # Enforce cache size limit
+    if len(XET_CONNECTION_INFO_CACHE) >= XET_CONNECTION_INFO_CACHE_SIZE:
+        XET_CONNECTION_INFO_CACHE.pop(next(iter(XET_CONNECTION_INFO_CACHE)))
+
+    # Update cache
+    XET_CONNECTION_INFO_CACHE[cache_key] = metadata
+
     return metadata
+
+
+def _cache_key(url: str, headers: dict[str, str], params: Optional[dict[str, str]]) -> str:
+    """Return a unique cache key for the given request parameters."""
+    lower_headers = {k.lower(): v for k, v in headers.items()}  # casing is not guaranteed here
+    auth_header = lower_headers.get("authorization", "")
+    params_str = "&".join(f"{k}={v}" for k, v in sorted((params or {}).items(), key=lambda x: x[0]))
+    return f"{url}|{auth_header}|{params_str}"
+
+
+def _is_expired(connection_info: XetConnectionInfo) -> bool:
+    """Check if the given XET connection info is expired."""
+    return connection_info.expiration_unix_epoch <= int(time.time()) + XET_CONNECTION_INFO_SAFETY_PERIOD