Repository: PrefectHQ/prefect

Issue:
Ability to customize retry status codes for client
### First check

- [X] I added a descriptive title to this issue.
- [X] I used the GitHub search to find a similar request and didn't find it.
- [X] I searched the Prefect documentation for this feature.

### Prefect Version

2.x

### Describe the current behavior

If the client fails to connect with the API for any reason other than 429, it fails, thus failing the flow. If there is an intermittent failure of network infrastructure between client and API server a retry can usually be successful mitigation. It would be helpful to be able to retry additional status codes via a setting.

### Describe the proposed behavior

Allow extras besides the 429 retry
https://github.com/PrefectHQ/prefect/blob/main/src/prefect/client.py#L260

via a setting (e.g. `PREFECT_API_ADDITIONAL_RETRY_HTTP_STATUS=500,502` 
https://github.com/PrefectHQ/prefect/blob/main/src/prefect/client.py#L260

Could also simplify the setting to `PREFECT_API_RETRY_ANY_FAILURE=true`


### Example Use

`prefect config set PREFECT_API_ADDITIONAL_RETRY_HTTP_STATUS=500,502` 
or 
`prefect config set PREFECT_API_RETRY_ANY_FAILURE=true`

### Additional context

_No response_


Discussion:
That implementation is specific to 429's in order to respect the rate limit back-off headers. If feasible, it'd be good to reuse standard retry handling for httpx. I'm not sure that exists though and our existing implementation _does_ have a fallback when retry-after is not specified. 

See https://github.com/encode/httpx/issues/108

We also may need to use a different retry mechanism for connections https://www.python-httpx.org/advanced/#usage_1

I'd recommend changing the prefix of these settings to `PREFECT_CLIENT`. For example:
 - `PREFECT_CLIENT_RETRY_STATUS_CODES=500,502`
 - `PREFECT_CLIENT_RETRY_METHODS=POST,PATCH`

We can define defaults for these in our settings and if you want to override them you must provide a comprehensive list. This allows users to opt out of retries for some codes/methods.
@abrookins it'd be great to see this land on the roadmap as retries are important for handling service disruptions.
it would be also useful to have a variable number of max retries rather than the hardcoded value https://github.com/PrefectHQ/prefect/blob/main/src/prefect/client.py#L252 e.g `PREFECT_CLIENT_MAX_RETRIES=5`
Adding a use case here: when deploying an API server behind an AWS load balancer, it's fairly common to get 502 errors (Bad Gateway) from the ALB. Retries would generally work great here. Related issue: https://github.com/PrefectHQ/prefect/issues/8080

I see the 429 and 503 error codes here: https://github.com/PrefectHQ/prefect/blame/30ecd3b711b9da07891393c55b5e28719a1505cf/src/prefect/client/base.py#L243

But the original solution proposed in this PR, to make this list customizable, would be better and would allow users to configure the retries to work with deployment-specific issues.
Happy to review pull requests adding `PREFECT_CLIENT_RETRY_STATUS_CODES` and `PREFECT_CLIENT_RETRY_COUNT`; preferably separately.

Resolving patch:
diff --git a/src/prefect/client/base.py b/src/prefect/client/base.py
--- a/src/prefect/client/base.py
+++ b/src/prefect/client/base.py
@@ -15,7 +15,10 @@
 
 from prefect.exceptions import PrefectHTTPStatusError
 from prefect.logging import get_logger
-from prefect.settings import PREFECT_CLIENT_RETRY_JITTER_FACTOR
+from prefect.settings import (
+    PREFECT_CLIENT_RETRY_EXTRA_CODES,
+    PREFECT_CLIENT_RETRY_JITTER_FACTOR,
+)
 from prefect.utilities.math import bounded_poisson_interval, clamped_poisson_interval
 
 # Datastores for lifespan management, keys should be a tuple of thread and app identities.
@@ -245,6 +248,7 @@ async def send(self, *args, **kwargs) -> Response:
             retry_codes={
                 status.HTTP_429_TOO_MANY_REQUESTS,
                 status.HTTP_503_SERVICE_UNAVAILABLE,
+                *PREFECT_CLIENT_RETRY_EXTRA_CODES.value(),
             },
             retry_exceptions=(
                 httpx.ReadTimeout,
diff --git a/src/prefect/settings.py b/src/prefect/settings.py
--- a/src/prefect/settings.py
+++ b/src/prefect/settings.py
@@ -277,6 +277,26 @@ def default_ui_api_url(settings, value):
     )(settings, value)
 
 
+def status_codes_as_integers_in_range(_, value):
+    """
+    `value_callback` for `PREFECT_CLIENT_RETRY_EXTRA_CODES` that ensures status codes
+    are integers in the range 100-599.
+    """
+    if value == "":
+        return set()
+
+    values = {v.strip() for v in value.split(",")}
+
+    if any(not v.isdigit() or int(v) < 100 or int(v) > 599 for v in values):
+        raise ValueError(
+            "PREFECT_CLIENT_RETRY_EXTRA_CODES must be a comma separated list of "
+            "integers between 100 and 599."
+        )
+
+    values = {int(v) for v in values}
+    return values
+
+
 def template_with_settings(*upstream_settings: Setting) -> Callable[["Settings", T], T]:
     """
     Returns a `value_callback` that will template the given settings into the runtime
@@ -525,6 +545,15 @@ def default_cloud_ui_url(settings, value):
 can affect retry lengths.
 """
 
+PREFECT_CLIENT_RETRY_EXTRA_CODES = Setting(
+    str, default="", value_callback=status_codes_as_integers_in_range
+)
+"""
+A comma-separated list of extra HTTP status codes to retry on. Defaults to an empty string.
+429 and 503 are always retried. Please note that not all routes are idempotent and retrying
+may result in unexpected behavior.
+"""
+
 PREFECT_CLOUD_API_URL = Setting(
     str,
     default="https://api.prefect.cloud/api",


Test patch:
diff --git a/tests/client/test_base_client.py b/tests/client/test_base_client.py
--- a/tests/client/test_base_client.py
+++ b/tests/client/test_base_client.py
@@ -7,7 +7,11 @@
 
 from prefect.client.base import PrefectHttpxClient, PrefectResponse
 from prefect.exceptions import PrefectHTTPStatusError
-from prefect.settings import PREFECT_CLIENT_RETRY_JITTER_FACTOR, temporary_settings
+from prefect.settings import (
+    PREFECT_CLIENT_RETRY_EXTRA_CODES,
+    PREFECT_CLIENT_RETRY_JITTER_FACTOR,
+    temporary_settings,
+)
 from prefect.testing.utilities import AsyncMock
 
 RESPONSE_429_RETRY_AFTER_0 = Response(
@@ -74,6 +78,72 @@ async def test_prefect_httpx_client_retries_on_designated_error_codes(
         assert "Another attempt will be made in 4s" in caplog.text
         assert "This is attempt 2/6" in caplog.text
 
+    @pytest.mark.usefixtures("mock_anyio_sleep", "disable_jitter")
+    @pytest.mark.parametrize(
+        "error_code,extra_codes",
+        [
+            (status.HTTP_408_REQUEST_TIMEOUT, "408"),
+            (status.HTTP_409_CONFLICT, "408,409"),
+        ],
+    )
+    async def test_prefect_httpx_client_retries_on_extra_error_codes(
+        self, monkeypatch, error_code, extra_codes, caplog
+    ):
+        base_client_send = AsyncMock()
+        monkeypatch.setattr(AsyncClient, "send", base_client_send)
+        client = PrefectHttpxClient()
+        retry_response = Response(
+            error_code,
+            request=Request("a test request", "fake.url/fake/route"),
+        )
+        base_client_send.side_effect = [
+            retry_response,
+            retry_response,
+            retry_response,
+            RESPONSE_200,
+        ]
+        with temporary_settings({PREFECT_CLIENT_RETRY_EXTRA_CODES: extra_codes}):
+            response = await client.post(
+                url="fake.url/fake/route", data={"evenmorefake": "data"}
+            )
+        assert response.status_code == status.HTTP_200_OK
+        assert base_client_send.call_count == 4
+
+        # We log on retry
+        assert "Received response with retryable status code" in caplog.text
+        assert "Another attempt will be made in 2s" in caplog.text
+        assert "This is attempt 1/6" in caplog.text
+
+        # A traceback should not be included
+        assert "Traceback" not in caplog.text
+
+        # Ensure the messaging changes
+        assert "Another attempt will be made in 4s" in caplog.text
+        assert "This is attempt 2/6" in caplog.text
+
+    @pytest.mark.usefixtures("mock_anyio_sleep", "disable_jitter")
+    async def test_prefect_httpx_client_raises_on_non_extra_error_codes(
+        self, monkeypatch, caplog
+    ):
+        base_client_send = AsyncMock()
+        monkeypatch.setattr(AsyncClient, "send", base_client_send)
+        client = PrefectHttpxClient()
+        retry_response = Response(
+            status.HTTP_408_REQUEST_TIMEOUT,
+            request=Request("a test request", "fake.url/fake/route"),
+        )
+        base_client_send.side_effect = [
+            retry_response,
+            retry_response,
+            retry_response,
+            RESPONSE_200,
+        ]
+        with temporary_settings({PREFECT_CLIENT_RETRY_EXTRA_CODES: "409"}):
+            with pytest.raises(PrefectHTTPStatusError):
+                await client.post(
+                    url="fake.url/fake/route", data={"evenmorefake": "data"}
+                )
+
     @pytest.mark.usefixtures("mock_anyio_sleep", "disable_jitter")
     @pytest.mark.parametrize(
         "exception_type",
diff --git a/tests/test_settings.py b/tests/test_settings.py
--- a/tests/test_settings.py
+++ b/tests/test_settings.py
@@ -14,6 +14,7 @@
     PREFECT_API_DATABASE_PASSWORD,
     PREFECT_API_KEY,
     PREFECT_API_URL,
+    PREFECT_CLIENT_RETRY_EXTRA_CODES,
     PREFECT_CLOUD_API_URL,
     PREFECT_CLOUD_UI_URL,
     PREFECT_CLOUD_URL,
@@ -473,6 +474,36 @@ def test_cloud_ui_url_set_directly(self):
         with temporary_settings({PREFECT_CLOUD_UI_URL: "test"}):
             assert PREFECT_CLOUD_UI_URL.value() == "test"
 
+    @pytest.mark.parametrize(
+        "extra_codes,expected",
+        [
+            ("", set()),
+            ("400", {400}),
+            ("400,400,400", {400}),
+            ("400,500", {400, 500}),
+            ("400, 401, 402", {400, 401, 402}),
+        ],
+    )
+    def test_client_retry_extra_codes(self, extra_codes, expected):
+        with temporary_settings({PREFECT_CLIENT_RETRY_EXTRA_CODES: extra_codes}):
+            assert PREFECT_CLIENT_RETRY_EXTRA_CODES.value() == expected
+
+    @pytest.mark.parametrize(
+        "extra_codes",
+        [
+            "foo",
+            "-1",
+            "0",
+            "10",
+            "400,foo",
+            "400,500,foo",
+        ],
+    )
+    def test_client_retry_extra_codes_invalid(self, extra_codes):
+        with pytest.raises(ValueError):
+            with temporary_settings({PREFECT_CLIENT_RETRY_EXTRA_CODES: extra_codes}):
+                PREFECT_CLIENT_RETRY_EXTRA_CODES.value()
+
 
 class TestTemporarySettings:
     def test_temporary_settings(self):