Voice et bot modif

venv/lib/python3.12/site-packages/huggingface_hub/utils/_xet.py

@@ -1,12 +1,13 @@
+import os
+import threading
 import time
 from dataclasses import dataclass
 from enum import Enum
-from typing import Optional
 
 import httpx
 
 from .. import constants
-from . import get_session, hf_raise_for_status, validate_hf_hub_args
+from . import hf_raise_for_status, http_backoff, validate_hf_hub_args
 
 
 XET_CONNECTION_INFO_SAFETY_PERIOD = 60  # seconds
@@ -32,9 +33,7 @@ class XetConnectionInfo:
     endpoint: str
 
 
-def parse_xet_file_data_from_response(
-    response: httpx.Response, endpoint: Optional[str] = None
-) -> Optional[XetFileData]:
+def parse_xet_file_data_from_response(response: httpx.Response, endpoint: str | None = None) -> XetFileData | None:
     """
     Parse XET file metadata from an HTTP response.
 
@@ -69,7 +68,7 @@ def parse_xet_file_data_from_response(
     )
 
 
-def parse_xet_connection_info_from_headers(headers: dict[str, str]) -> Optional[XetConnectionInfo]:
+def parse_xet_connection_info_from_headers(headers: dict[str, str]) -> XetConnectionInfo | None:
     """
     Parse XET connection info from the HTTP headers or return None if not found.
     Args:
@@ -122,16 +121,53 @@ def refresh_xet_connection_info(
     return _fetch_xet_connection_info_with_url(file_data.refresh_route, headers)
 
 
+@validate_hf_hub_args
+def xet_connection_info_refresh_url(
+    *,
+    token_type: XetTokenType,
+    repo_id: str,
+    repo_type: str,
+    revision: str | None = None,
+    endpoint: str | None = None,
+) -> str:
+    """
+    Build the URL used to fetch or refresh a Xet access token for a given repo.
+    Args:
+        token_type (`XetTokenType`):
+            Type of the token to request: `"read"` or `"write"`.
+        repo_id (`str`):
+            A namespace (user or an organization) and a repo name separated by a `/`.
+        repo_type (`str`):
+            Type of the repo (e.g. `"model"`, `"dataset"`, `"space"`, `"bucket"`).
+        revision (`str`, `optional`):
+            The revision of the repo to get the token for.
+        endpoint (`str`, `optional`):
+            The endpoint to use for the request. Defaults to the Hub endpoint.
+    Returns:
+        `str`:
+            The fully-qualified URL of the token refresh endpoint.
+    """
+    endpoint = endpoint if endpoint is not None else constants.ENDPOINT
+    url = f"{endpoint}/api/{repo_type}s/{repo_id}/xet-{token_type.value}-token"
+    if repo_type != "bucket" or revision is not None:
+        # On "bucket" repo type, the revision never needed => don't use it
+        # Otherwise, use the revision.
+        # Note: when creating a PR on a git-based repo, user needs write access but they don't know the revision in advance.
+        # => pass "/None" in URL and server will return a token for PR refs.
+        url += f"/{revision}"
+    return url
+
+
 @validate_hf_hub_args
 def fetch_xet_connection_info_from_repo_info(
     *,
     token_type: XetTokenType,
     repo_id: str,
     repo_type: str,
-    revision: Optional[str] = None,
+    revision: str | None = None,
     headers: dict[str, str],
-    endpoint: Optional[str] = None,
-    params: Optional[dict[str, str]] = None,
+    endpoint: str | None = None,
+    params: dict[str, str] | None = None,
 ) -> XetConnectionInfo:
     """
     Uses the repo info to request a xet access token from Hub.
@@ -159,16 +195,22 @@ def fetch_xet_connection_info_from_repo_info(
         [`ValueError`](https://docs.python.org/3/library/exceptions.html#ValueError)
             If the Hub API response is improperly formatted.
     """
-    endpoint = endpoint if endpoint is not None else constants.ENDPOINT
-    url = f"{endpoint}/api/{repo_type}s/{repo_id}/xet-{token_type.value}-token/{revision}"
-    return _fetch_xet_connection_info_with_url(url, headers, params)
+    url = xet_connection_info_refresh_url(
+        token_type=token_type,
+        repo_id=repo_id,
+        repo_type=repo_type,
+        revision=revision,
+        endpoint=endpoint,
+    )
+    return _fetch_xet_connection_info_with_url(url, headers, params, cache_key_prefix=f"{repo_type}-{repo_id}")
 
 
 @validate_hf_hub_args
 def _fetch_xet_connection_info_with_url(
     url: str,
     headers: dict[str, str],
-    params: Optional[dict[str, str]] = None,
+    params: dict[str, str] | None = None,
+    cache_key_prefix: str | None = None,
 ) -> XetConnectionInfo:
     """
     Requests the xet connection info from the supplied URL. This includes the
@@ -193,14 +235,14 @@ def _fetch_xet_connection_info_with_url(
             If the Hub API response is improperly formatted.
     """
     # Check cache first
-    cache_key = _cache_key(url, headers, params)
+    cache_key = _cache_key(url, headers, params, prefix=cache_key_prefix)
     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)
+    resp = http_backoff("GET", url, headers=headers, params=params)
     hf_raise_for_status(resp)
 
     metadata = parse_xet_connection_info_from_headers(resp.headers)  # type: ignore
@@ -222,14 +264,107 @@ def _fetch_xet_connection_info_with_url(
     return metadata
 
 
-def _cache_key(url: str, headers: dict[str, str], params: Optional[dict[str, str]]) -> str:
+def reset_xet_connection_info_cache_for_repo(repo_type: str | None, repo_id: str) -> None:
+    """Reset the XET connection info cache for the given repo type and repo id.
+
+    Used when a repo is deleted.
+    """
+    if repo_type is None:
+        repo_type = constants.REPO_TYPE_MODEL
+    prefix = f"{repo_type}-{repo_id}|"
+    for k in list(XET_CONNECTION_INFO_CACHE.keys()):
+        if k.startswith(prefix):
+            XET_CONNECTION_INFO_CACHE.pop(k, None)
+
+
+def _cache_key(url: str, headers: dict[str, str], params: dict[str, str] | None, prefix: str | None = None) -> 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}"
+    return f"{prefix}|{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
+
+
+class XetSessionHolder:
+    """Holds an optional XetSession; supports safe re-creation after sigint_abort or fork.
+
+    Thread-safe: a ``threading.Lock`` guards all state mutations, which matters
+    for free-threaded Python (3.14t) where multiple threads can race on ``get()``
+    or ``sigint_abort()`` without the GIL serialising them.
+    """
+
+    def __init__(self):
+        self._lock = threading.Lock()
+        self._session = None
+        self._session_pid: int | None = None
+
+    def get(self):
+        """Return the current session, creating one if needed.
+
+        Fork-safe: if the current process PID differs from the PID that created
+        the session (i.e. we are in a forked child), the old session is discarded
+        and a fresh session is created for this process.
+        """
+        with self._lock:
+            current_pid = os.getpid()
+
+            if self._session is not None and self._session_pid != current_pid:
+                # Fork detected. Discard the parent's session; the Rust Drop will
+                # call discard_runtime() (std::mem::forget) rather than the normal
+                # shutdown path, so this returns immediately without blocking.
+                self._session = None
+
+            if self._session is None:
+                from hf_xet import XetSession
+
+                self._session = XetSession()
+                self._session_pid = current_pid
+
+            return self._session
+
+    def sigint_abort(self):
+        """Abort the current session and clear it so the next get() creates a fresh one."""
+        with self._lock:
+            if self._session is not None:
+                try:
+                    self._session.sigint_abort()
+                except Exception:
+                    pass
+                self._session = None
+                self._session_pid = None
+
+
+_GLOBAL_XET_HOLDER = XetSessionHolder()
+
+
+def get_xet_session():
+    """Return the global :class:`hf_xet.XetSession`, creating it on first call.
+
+    The session is shared across all calls within a process, just as the HTTP
+    client returned by :func:`~huggingface_hub.utils._http.get_session` is shared.
+    It is created lazily and is fork-safe and thread-safe.
+    """
+    return _GLOBAL_XET_HOLDER.get()
+
+
+def xet_headers_without_auth(headers: dict[str, str]) -> dict[str, str]:
+    """Return a copy of headers with the authorization header removed.
+
+    Xet storage requests use a short-lived xet access token for auth, so the
+    Hub authorization header must not be forwarded to xet storage endpoints.
+    """
+    return {key: value for key, value in headers.items() if key.lower() != "authorization"}
+
+
+def abort_xet_session():
+    """Abort the global xet session after a KeyboardInterrupt.
+
+    Cancels any in-flight Rust operation and clears the session so the next
+    call to :func:`get_xet_session` starts fresh (notebook-friendly).
+    """
+    _GLOBAL_XET_HOLDER.sigint_abort()