Add request ID to all requests (#909)

* Add request ID to all requests

* Typo

* Typo

* Review comments

* Update src/huggingface_hub/utils/_errors.py

Co-authored-by: Sylvain Gugger <[email protected]>

Co-authored-by: Sylvain Gugger <[email protected]>

Author: LysandreJik

src/huggingface_hub/commands/lfs.py

@@ -29,6 +29,7 @@
 from huggingface_hub.lfs import LFS_MULTIPART_UPLOAD_COMMAND
 
 from ..utils import logging
+from ..utils._errors import _raise_with_request_id
 
 
 logger = logging.get_logger(__name__)
@@ -212,7 +213,7 @@ def run(self):
                     filepath, seek_from=i * chunk_size, read_limit=chunk_size
                 ) as data:
                     r = requests.put(presigned_url, data=data)
-                    r.raise_for_status()
+                    _raise_with_request_id(r)
                     parts.append(
                         {
                             "etag": r.headers.get("etag"),
@@ -238,6 +239,6 @@ def run(self):
                     "parts": parts,
                 },
             )
-            r.raise_for_status()
+            _raise_with_request_id(r)
 
             write_msg({"event": "complete", "oid": oid})

src/huggingface_hub/hf_api.py

@@ -32,7 +32,7 @@
     SPACES_SDK_TYPES,
 )
 from .utils import logging
-from .utils._errors import _raise_for_status
+from .utils._errors import _raise_for_status, _raise_with_request_id
 from .utils._fixes import JSONDecodeError
 from .utils.endpoint_helpers import (
     AttributeDictionary,
@@ -535,7 +535,7 @@ def whoami(self, token: Optional[str] = None) -> Dict:
         path = f"{self.endpoint}/api/whoami-v2"
         r = requests.get(path, headers={"authorization": f"Bearer {token}"})
         try:
-            r.raise_for_status()
+            _raise_with_request_id(r)
         except HTTPError as e:
             raise HTTPError(
                 "Invalid user token. If you didn't pass a user token, make sure you "
@@ -629,7 +629,7 @@ def get_model_tags(self) -> ModelTags:
         "Gets all valid model tags as a nested namespace object"
         path = f"{self.endpoint}/api/models-tags-by-type"
         r = requests.get(path)
-        r.raise_for_status()
+        _raise_with_request_id(r)
         d = r.json()
         return ModelTags(d)
 
@@ -639,7 +639,7 @@ def get_dataset_tags(self) -> DatasetTags:
         """
         path = f"{self.endpoint}/api/datasets-tags-by-type"
         r = requests.get(path)
-        r.raise_for_status()
+        _raise_with_request_id(r)
         d = r.json()
         return DatasetTags(d)
 
@@ -778,7 +778,7 @@ def list_models(
         if cardData is not None:
             params.update({"cardData": cardData})
         r = requests.get(path, params=params, headers=headers)
-        r.raise_for_status()
+        _raise_with_request_id(r)
         d = r.json()
         res = [ModelInfo(**x) for x in d]
         if emissions_thresholds is not None:
@@ -971,7 +971,7 @@ def list_datasets(
             if cardData:
                 params.update({"full": True})
         r = requests.get(path, params=params, headers=headers)
-        r.raise_for_status()
+        _raise_with_request_id(r)
         d = r.json()
         return [DatasetInfo(**x) for x in d]
 
@@ -1026,7 +1026,7 @@ def list_metrics(self) -> List[MetricInfo]:
         path = f"{self.endpoint}/api/metrics"
         params = {}
         r = requests.get(path, params=params)
-        r.raise_for_status()
+        _raise_with_request_id(r)
         d = r.json()
         return [MetricInfo(**x) for x in d]
 
@@ -1405,7 +1405,7 @@ def create_repo(
         )
 
         try:
-            r.raise_for_status()
+            _raise_with_request_id(r)
         except HTTPError as err:
             if not (exist_ok and err.response.status_code == 409):
                 try:

src/huggingface_hub/utils/_errors.py

@@ -52,27 +52,44 @@ def _raise_for_status(request):
     Internal version of `request.raise_for_status()` that will refine a
     potential HTTPError.
     """
+    request_id = request.headers.get("X-Request-Id")
+
     if "X-Error-Code" in request.headers:
         error_code = request.headers["X-Error-Code"]
         if error_code == "RepoNotFound":
             raise RepositoryNotFoundError(
-                f"404 Client Error: Repository Not Found for url: {request.url}. "
-                "If the repo is private, make sure you are authenticated."
+                f"404 Client Error: Repository Not Found for url: {request.url}. If the"
+                " repo is private, make sure you are authenticated. (Request ID:"
+                f" {request_id})"
             )
         elif error_code == "RevisionNotFound":
             raise RevisionNotFoundError(
-                f"404 Client Error: Revision Not Found for url: {request.url}"
+                f"404 Client Error: Revision Not Found for url: {request.url}. (Request"
+                f" ID: {request_id})"
             )
         elif error_code == "EntryNotFound":
             raise EntryNotFoundError(
-                f"404 Client Error: Entry Not Found for url: {request.url}"
+                f"404 Client Error: Entry Not Found for url: {request.url}. (Request"
+                f" ID: {request_id})"
             )
 
     if request.status_code == 401:
         # The repo was not found and the user is not Authenticated
         raise RepositoryNotFoundError(
-            f"401 Client Error: Repository Not Found for url: {request.url}. "
-            "If the repo is private, make sure you are authenticated."
+            f"401 Client Error: Repository Not Found for url: {request.url}. If the"
+            " repo is private, make sure you are authenticated. (Request ID:"
+            f" {request_id})"
         )
 
-    request.raise_for_status()
+    _raise_with_request_id(request)
+
+
+def _raise_with_request_id(request):
+    request_id = request.headers.get("X-Request-Id")
+    try:
+        request.raise_for_status()
+    except Exception as e:
+        if request_id is not None and len(e.args) > 0 and isinstance(e.args[0], str):
+            e.args = (e.args[0] + f" (Request ID: {request_id})",) + e.args[1:]
+
+        raise e

tests/test_hf_api.py

@@ -197,9 +197,9 @@ def test_whoami(self):
     @retry_endpoint
     def test_delete_repo_error_message(self):
         # test for #751
-        with pytest.raises(
-            HTTPError,
-            match="404 Client Error: Repository Not Found",
+        with self.assertRaisesRegex(
+            requests.exceptions.HTTPError,
+            r"404 Client Error: Repository Not Found (.+) \(Request ID: .+\)",
         ):
             self._api.delete_repo("repo-that-does-not-exist", token=self._token)
 
@@ -970,7 +970,8 @@ def test_model_info(self):
         shutil.rmtree(os.path.dirname(HfFolder.path_token))
         # Test we cannot access model info without a token
         with self.assertRaisesRegex(
-            requests.exceptions.HTTPError, "401 Client Error: Repository Not Found"
+            requests.exceptions.HTTPError,
+            r"401 Client Error: Repository Not Found for url: (.+) \(Request ID: .+\)",
         ):
             _ = self._api.model_info(repo_id=f"{USER}/{self.REPO_NAME}")
         # Test we can access model info with a token