Handle Rate Limits in Pagination with Automatic Retries (#2970)

* rate limit handling

* Import http_backoff (missed it earlier, mb)

* fix tests

---------

Co-authored-by: Lucain <[email protected]>
Co-authored-by: Lucain Pouget <[email protected]>

Author: Weyaxi

src/huggingface_hub/utils/_pagination.py

@@ -18,7 +18,7 @@
 
 import requests
 
-from . import get_session, hf_raise_for_status, logging
+from . import get_session, hf_raise_for_status, http_backoff, logging
 
 
 logger = logging.get_logger(__name__)
@@ -42,7 +42,7 @@ def paginate(path: str, params: Dict, headers: Dict) -> Iterable:
     next_page = _get_next_page(r)
     while next_page is not None:
         logger.debug(f"Pagination detected. Requesting next page: {next_page}")
-        r = session.get(next_page, headers=headers)
+        r = http_backoff("GET", next_page, max_retries=20, retry_on_status_codes=429, headers=headers)
         hf_raise_for_status(r)
         yield from r.json()
         next_page = _get_next_page(r)

tests/test_utils_pagination.py

@@ -8,9 +8,12 @@
 
 class TestPagination(unittest.TestCase):
     @patch("huggingface_hub.utils._pagination.get_session")
+    @patch("huggingface_hub.utils._pagination.http_backoff")
     @patch("huggingface_hub.utils._pagination.hf_raise_for_status")
     @handle_injection_in_test
-    def test_mocked_paginate(self, mock_get_session: Mock, mock_hf_raise_for_status: Mock) -> None:
+    def test_mocked_paginate(
+        self, mock_get_session: Mock, mock_http_backoff: Mock, mock_hf_raise_for_status: Mock
+    ) -> None:
         mock_get = mock_get_session().get
         mock_params = Mock()
         mock_headers = Mock()
@@ -33,31 +36,32 @@ def test_mocked_paginate(self, mock_get_session: Mock, mock_hf_raise_for_status:
         # Mock response
         mock_get.side_effect = [
             mock_response_page_1,
+        ]
+        mock_http_backoff.side_effect = [
             mock_response_page_2,
             mock_response_page_3,
         ]
 
         results = paginate("url", params=mock_params, headers=mock_headers)
 
         # Requests are made only when generator is yielded
-        self.assertEqual(mock_get.call_count, 0)
+        assert mock_get.call_count == 0
 
         # Results after concatenating pages
-        self.assertListEqual(list(results), [1, 2, 3, 4, 5, 6, 7, 8])
+        assert list(results) == [1, 2, 3, 4, 5, 6, 7, 8]
 
         # All pages requested: 3 requests, 3 raise for status
-        self.assertEqual(mock_get.call_count, 3)
-        self.assertEqual(mock_hf_raise_for_status.call_count, 3)
+        # First request with `get_session.get` (we want at least 1 request to succeed correctly) and 2 with `http_backoff`
+        assert mock_get.call_count == 1
+        assert mock_http_backoff.call_count == 2
+        assert mock_hf_raise_for_status.call_count == 3
 
         # Params not passed to next pages
-        self.assertListEqual(
-            mock_get.call_args_list,
-            [
-                call("url", params=mock_params, headers=mock_headers),
-                call("url_p2", headers=mock_headers),
-                call("url_p3", headers=mock_headers),
-            ],
-        )
+        assert mock_get.call_args_list == [call("url", params=mock_params, headers=mock_headers)]
+        assert mock_http_backoff.call_args_list == [
+            call("GET", "url_p2", max_retries=20, retry_on_status_codes=429, headers=mock_headers),
+            call("GET", "url_p3", max_retries=20, retry_on_status_codes=429, headers=mock_headers),
+        ]
 
     def test_paginate_github_api(self) -> None:
         # Real test: paginate over huggingface repos on Github