Better streaming retries (504 and 429) (#7847)

better streaming retries

Author: lhoestq

src/datasets/config.py

@@ -248,6 +248,8 @@
 # Streaming
 STREAMING_READ_MAX_RETRIES = 20
 STREAMING_READ_RETRY_INTERVAL = 5
+STREAMING_READ_SERVER_UNAVAILABLE_RETRY_INTERVAL = 20
+STREAMING_READ_RATE_LIMIT_RETRY_INTERVAL = 60
 STREAMING_OPEN_MAX_RETRIES = 20
 STREAMING_OPEN_RETRY_INTERVAL = 5
 

src/datasets/utils/file_utils.py

@@ -68,6 +68,8 @@ class _AiohttpClientError(Exception):
     requests.exceptions.Timeout,
     httpx.RequestError,
 )
+SERVER_UNAVAILABLE_CODE = 504
+RATE_LIMIT_CODE = 429
 
 
 def is_remote_url(url_or_filename: str) -> bool:
@@ -827,6 +829,22 @@ def read_with_retries(*args, **kwargs):
                     f"Got disconnected from remote data host. Retrying in {config.STREAMING_READ_RETRY_INTERVAL}sec [{retry}/{max_retries}]"
                 )
                 time.sleep(config.STREAMING_READ_RETRY_INTERVAL)
+            except huggingface_hub.errors.HfHubHTTPError as err:
+                if err.response is not None and err.response.status_code == SERVER_UNAVAILABLE_CODE:
+                    disconnect_err = err
+                    logger.warning(
+                        f"Got disconnected from remote data host. Retrying in {config.STREAMING_READ_SERVER_UNAVAILABLE_RETRY_INTERVAL}sec [{retry}/{max_retries}]"
+                    )
+                    time.sleep(config.STREAMING_READ_SERVER_UNAVAILABLE_RETRY_INTERVAL)
+                elif err.response is not None and err.response.status_code == RATE_LIMIT_CODE:
+                    disconnect_err = err
+                    logger.warning(str(err))
+                    logger.warning(
+                        f"Got disconnected from remote data host. Retrying in {config.STREAMING_READ_RATE_LIMIT_RETRY_INTERVAL}sec [{retry}/{max_retries}]"
+                    )
+                    time.sleep(config.STREAMING_READ_RATE_LIMIT_RETRY_INTERVAL)
+                else:
+                    raise
         else:
             raise ConnectionError("Server Disconnected") from disconnect_err
         return out