Automatically retry on 429 Too Many Requests

soxofaan · soxofaan · commit 964441e384a6 · 2025-05-09T17:27:30.000+02:00
refs #441, #764
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -11,6 +11,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - Support `collection_property` based property filtering in `load_stac` ([#246](https://github.com/Open-EO/openeo-python-client/issues/246))
 - Add `validate()` method to `SaveResult`, `VectorCube`, `MlModel` and `StacResource` classes ([#766](https://github.com/Open-EO/openeo-python-client/issues/766))
+- Automatically retry on `429 Too Many Requests` HTTP errors ([#441](https://github.com/Open-EO/openeo-python-client/issues/441))
 
 ### Changed
 
diff --git a/openeo/extra/job_management/__init__.py b/openeo/extra/job_management/__init__.py
@@ -284,7 +284,7 @@ def _make_resilient(connection):
         503 Service Unavailable
         504 Gateway Timeout
         """
-        # TODO: refactor this helper out of this class and unify with `openeo_driver.util.http.requests_with_retry`
+        # TODO: migrate this to now built-in retry configuration of `Connection` or `openeo.util.http.retry_adapter`?
         status_forcelist = [500, 502, 503, 504]
         retries = Retry(
             total=MAX_RETRIES,
diff --git a/openeo/rest/_connection.py b/openeo/rest/_connection.py
@@ -5,13 +5,15 @@
 from typing import Iterable, Optional, Union
 
 import requests
+import requests.adapters
 from requests import Response
 from requests.auth import AuthBase
 
 import openeo
 from openeo.rest import OpenEoApiError, OpenEoApiPlainError, OpenEoRestError
 from openeo.rest.auth.auth import NullAuth
 from openeo.util import ContextTimer, ensure_list, str_truncate, url_join
+from openeo.utils.http import session_with_retries
 
 _log = logging.getLogger(__name__)
 
@@ -26,15 +28,22 @@ class RestApiConnection:
     def __init__(
         self,
         root_url: str,
+        *,
         auth: Optional[AuthBase] = None,
         session: Optional[requests.Session] = None,
         default_timeout: Optional[int] = None,
         slow_response_threshold: Optional[float] = None,
+        retry: Union[requests.adapters.Retry, dict, bool, None] = None,
     ):
         self._root_url = root_url
         self._auth = None
         self.auth = auth or NullAuth()
-        self.session = session or requests.Session()
+        if session:
+            self.session = session
+        elif retry is not False:
+            self.session = session_with_retries(retry=retry)
+        else:
+            self.session = requests.Session()
         self.default_timeout = default_timeout or DEFAULT_TIMEOUT
         self.default_headers = {
             "User-Agent": "openeo-python-client/{cv} {py}/{pv} {pl}".format(
diff --git a/openeo/rest/connection.py b/openeo/rest/connection.py
@@ -28,6 +28,7 @@
 )
 
 import requests
+import requests.adapters
 import shapely.geometry.base
 from requests.auth import AuthBase, HTTPBasicAuth
 
@@ -114,6 +115,16 @@ class Connection(RestApiConnection):
         optional :class:`OidcAuthenticator` object to use for renewing OIDC tokens.
     :param auth: Optional ``requests.auth.AuthBase`` object to use for requests.
         Usage of this parameter is deprecated, use the specific authentication methods instead.
+    :param retry: general request retry settings, can be specified as:
+
+        - :py:class:`requests.adapters.Retry` object
+          or a dictionary with corresponding keyword arguments
+          (e.g. ``total``, ``backoff_factor``, ``status_forcelist``, ...)
+        - ``None`` (default) to use default openEO-oriented retry settings
+        - ``False`` to disable retrying requests
+
+    .. versionchanged:: 0.41.0
+        Added ``retry`` argument.
     """
 
     _MINIMUM_API_VERSION = ComparableVersion("1.0.0")
@@ -130,6 +141,7 @@ def __init__(
         refresh_token_store: Optional[RefreshTokenStore] = None,
         oidc_auth_renewer: Optional[OidcAuthenticator] = None,
         auth: Optional[AuthBase] = None,
+        retry: Union[requests.adapters.Retry, dict, bool, None] = None,
     ):
         if "://" not in url:
             url = "https://" + url
@@ -139,6 +151,7 @@ def __init__(
             root_url=self.version_discovery(url, session=session, timeout=default_timeout),
             auth=auth, session=session, default_timeout=default_timeout,
             slow_response_threshold=slow_response_threshold,
+            retry=retry,
         )
 
         # Initial API version check.
@@ -1885,6 +1898,7 @@ def connect(
     session: Optional[requests.Session] = None,
     default_timeout: Optional[int] = None,
     auto_validate: bool = True,
+    retry: Union[requests.adapters.Retry, dict, bool, None] = None,
 ) -> Connection:
     """
     This method is the entry point to OpenEO.
@@ -1904,9 +1918,19 @@ def connect(
     :param auth_options: Options/arguments specific to the authentication type
     :param default_timeout: default timeout (in seconds) for requests
     :param auto_validate: toggle to automatically validate process graphs before execution
+    :param retry: general request retry settings, can be specified as:
+
+        - :py:class:`requests.adapters.Retry` object
+          or a dictionary with corresponding keyword arguments
+          (e.g. ``total``, ``backoff_factor``, ``status_forcelist``, ...)
+        - ``None`` (default) to use default openEO-oriented retry settings
+        - ``False`` to disable retrying requests
 
-    .. versionadded:: 0.24.0
-        added ``auto_validate`` argument
+    .. versionchanged:: 0.24.0
+        Added ``auto_validate`` argument
+
+    .. versionchanged:: 0.41.0
+        Added ``retry`` argument.
     """
 
     def _config_log(message):
@@ -1931,7 +1955,9 @@ def _config_log(message):
 
     if not url:
         raise OpenEoClientException("No openEO back-end URL given or known to connect to.")
-    connection = Connection(url, session=session, default_timeout=default_timeout, auto_validate=auto_validate)
+    connection = Connection(
+        url, session=session, default_timeout=default_timeout, auto_validate=auto_validate, retry=retry
+    )
 
     auth_type = auth_type.lower() if isinstance(auth_type, str) else auth_type
     if auth_type in {None, False, 'null', 'none'}:
diff --git a/openeo/utils/http.py b/openeo/utils/http.py
@@ -0,0 +1,96 @@
+"""
+openEO-oriented HTTP utilities
+"""
+
+from typing import Collection, Union
+
+import requests
+import requests.adapters
+
+DEFAULT_RETRIES_TOTAL = 5
+
+# On `backoff_factor`: it influences how much to sleep according to the formula:
+#     sleep = {backoff factor} * (2 ** ({consecutive errors - 1}))
+# The sleep before the first retry will be skipped however.
+# For example with backoff_factor=2.5, the sleeps between consecutive attempts would be:
+#     0, 5, 10, 20, 40, ...
+DEFAULT_BACKOFF_FACTOR = 2.5
+
+
+DEFAULT_RETRY_FORCELIST = frozenset(
+    [
+        429,  # Too Many Requests
+        500,  # Internal Server Error
+        502,  # Bad Gateway
+        503,  # Service Unavailable
+        504,  # Gateway Timeout
+    ]
+)
+
+
+def retry_adapter(
+    *,
+    total: int = DEFAULT_RETRIES_TOTAL,
+    backoff_factor: float = DEFAULT_BACKOFF_FACTOR,
+    status_forcelist: Collection[int] = DEFAULT_RETRY_FORCELIST,
+    **kwargs,
+) -> requests.adapters.Retry:
+    """
+    Factory for creating a `requests.adapters.Retry` configuration object with
+    openEO-oriented retry settings.
+
+    :param total: Total number of retries to allow
+    :param backoff_factor: scaling factor for sleeps between retries
+    :param status_forcelist: A set of integer HTTP status codes that we should force a retry on.
+    :param kwargs: additional kwargs to pass to `requests.adapters.Retry`
+    :return:
+
+    Inspiration and references:
+    - https://requests.readthedocs.io/en/latest/api/#requests.adapters.HTTPAdapter
+    - https://urllib3.readthedocs.io/en/latest/reference/urllib3.util.html#urllib3.util.Retry
+    - https://findwork.dev/blog/advanced-usage-python-requests-timeouts-retries-hooks/#retry-on-failure
+    """
+    retry = requests.adapters.Retry(
+        total=total,
+        backoff_factor=backoff_factor,
+        status_forcelist=status_forcelist,
+        **kwargs,
+    )
+    return retry
+
+
+def _to_retry(
+    retry: Union[requests.adapters.Retry, dict, None],
+) -> requests.adapters.Retry:
+    """
+    Convert a retry specification to a `requests.adapters.Retry` object.
+    """
+    if isinstance(retry, requests.adapters.Retry):
+        return retry
+    elif isinstance(retry, dict):
+        adapter = retry_adapter(**retry)
+    elif retry in {None, True}:
+        adapter = retry_adapter()
+    else:
+        raise ValueError(f"Invalid retry setting: {retry!r}")
+    return adapter
+
+
+def session_with_retries(
+    retry: Union[requests.adapters.Retry, dict, None] = None,
+) -> requests.Session:
+    """
+    Factory for a requests session with openEO-oriented retry settings.
+
+    :param retry: The retry configuration, can be specified as:
+        - :py:class:`requests.adapters.Retry`
+        - a dictionary with :py:class:`requests.adapters.Retry` arguments,
+          e.g. ``total``, ``backoff_factor``, ``status_forcelist``, ...
+        - ``None`` for default openEO-oriented retry settings
+    """
+    session = requests.Session()
+    retry = _to_retry(retry)
+    adapter = requests.adapters.HTTPAdapter(max_retries=retry)
+    session.mount("http://", adapter)
+    session.mount("https://", adapter)
+    return session
diff --git a/tests/rest/test_job.py b/tests/rest/test_job.py
@@ -1,3 +1,4 @@
+import contextlib
 import itertools
 import json
 import logging
@@ -6,12 +7,19 @@
 from typing import Optional
 from unittest import mock
 
+import dirty_equals
+import httpretty
 import pytest
 import requests
 
 import openeo
 import openeo.rest.job
-from openeo.rest import JobFailedException, OpenEoApiPlainError, OpenEoClientException
+from openeo.rest import (
+    DEFAULT_JOB_STATUS_POLL_CONNECTION_RETRY_INTERVAL,
+    JobFailedException,
+    OpenEoApiPlainError,
+    OpenEoClientException,
+)
 from openeo.rest.job import BatchJob, ResultAsset
 from openeo.rest.models.general import Link
 from openeo.rest.models.logs import LogEntry
@@ -310,6 +318,90 @@ def test_execute_batch_with_excessive_soft_errors(con100, requests_mock, tmpdir,
     ]
 
 
+@httpretty.activate(allow_net_connect=False)
+@pytest.mark.parametrize(
+    ["retry", "expectation_context", "expected_sleeps"],
+    [
+        (  # Default retry settings
+            None,
+            contextlib.nullcontext(),
+            [0.1, 23, 34],
+        ),
+        (
+            # Only retry on 429 (and fail on 500)
+            {"status_forcelist": [429]},
+            pytest.raises(OpenEoApiPlainError, match=re.escape("[500] Internal Server Error")),
+            [0.1, 23, DEFAULT_JOB_STATUS_POLL_CONNECTION_RETRY_INTERVAL],
+        ),
+        (
+            # No retry setup
+            False,
+            pytest.raises(OpenEoApiPlainError, match=re.escape("[429] Too Many Requests")),
+            [0.1],
+        ),
+    ],
+)
+def test_execute_batch_retry_after_429_too_many_requests(tmpdir, retry, expectation_context, expected_sleeps):
+    httpretty.register_uri(
+        httpretty.GET,
+        uri=API_URL + "/",
+        body=json.dumps({"api_version": "1.0.0", "endpoints": [{"path": "/credentials/basic", "methods": ["GET"]}]}),
+    )
+    httpretty.register_uri(
+        httpretty.GET,
+        uri=API_URL + "/file_formats",
+        body=json.dumps({"output": {"GTiff": {"gis_data_types": ["raster"]}}}),
+    )
+    httpretty.register_uri(
+        httpretty.GET,
+        uri=API_URL + "/collections/SENTINEL2",
+        body=json.dumps({"foo": "bar"}),
+    )
+    httpretty.register_uri(
+        httpretty.POST, uri=API_URL + "/jobs", status=201, adding_headers={"OpenEO-Identifier": "f00ba5"}, body=""
+    )
+    httpretty.register_uri(httpretty.POST, uri=API_URL + "/jobs/f00ba5/results", status=202)
+    httpretty.register_uri(
+        httpretty.GET,
+        uri=API_URL + "/jobs/f00ba5",
+        responses=[
+            httpretty.Response(body=json.dumps({"status": "queued"})),
+            httpretty.Response(status=429, body="Too Many Requests", adding_headers={"Retry-After": "23"}),
+            httpretty.Response(body=json.dumps({"status": "running", "progress": 80})),
+            httpretty.Response(status=502, body="Bad Gateway"),
+            httpretty.Response(status=500, body="Internal Server Error"),
+            httpretty.Response(body=json.dumps({"status": "running", "progress": 80})),
+            httpretty.Response(status=429, body="Too Many Requests", adding_headers={"Retry-After": "34"}),
+            httpretty.Response(body=json.dumps({"status": "finished", "progress": 100})),
+        ],
+    )
+    httpretty.register_uri(
+        httpretty.GET,
+        uri=API_URL + "/jobs/f00ba5/results",
+        body=json.dumps(
+            {
+                "assets": {
+                    "output.tiff": {
+                        "href": API_URL + "/jobs/f00ba5/files/output.tiff",
+                        "type": "image/tiff; application=geotiff",
+                    },
+                }
+            }
+        ),
+    )
+    httpretty.register_uri(httpretty.GET, uri=API_URL + "/jobs/f00ba5/files/output.tiff", body="tiffdata")
+    httpretty.register_uri(httpretty.GET, uri=API_URL + "/jobs/f00ba5/logs", body=json.dumps({"logs": []}))
+
+    con = openeo.connect(API_URL, retry=retry)
+
+    with mock.patch("time.sleep") as sleep_mock:
+        job = con.load_collection("SENTINEL2").create_job()
+        with expectation_context:
+            job.start_and_wait(max_poll_interval=0.1)
+
+    assert sleep_mock.call_args_list == dirty_equals.Contains(*(mock.call(s) for s in expected_sleeps))
+
+
 class LogGenerator:
     """Helper to generate log entry (dicts) with auto-generated ids, messages, etc."""
 
diff --git a/tests/utils/test_http.py b/tests/utils/test_http.py
