diff --git a/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/adapter.py b/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/adapter.py index d67180020d..dd177bd7dc 100644 --- a/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/adapter.py +++ b/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/adapter.py @@ -61,7 +61,13 @@ def client_from_platform( _skip = {"accept", "accept-encoding", "connection", "user-agent", "host"} headers = {k: v for k, v in platform._client.headers.items() if k.lower() not in _skip} # type: ignore[union-attr] - retry = RetryPolicy(max_retries=platform.max_retries) + retry = RetryPolicy( + max_retries=platform.max_retries, + retryable_status_codes=(408, 409, 429), + retry_all_server_errors=True, + respect_retry_decision_headers=True, + respect_retry_after_headers=True, + ) url_resolver = _url_resolver_from_platform(platform) if isinstance(platform, AsyncNeMoPlatform): if not issubclass(client_cls, AsyncNemoClient): diff --git a/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/client.py b/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/client.py index feb69df85c..e42996a31e 100644 --- a/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/client.py +++ b/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/client.py @@ -18,12 +18,14 @@ import asyncio import copy +import email.utils import inspect import json import os import time from collections.abc import AsyncIterator, Callable, Iterator, Mapping from contextlib import asynccontextmanager, contextmanager +from datetime import timezone from functools import cache from pathlib import Path from typing import Any, Self, TypeVar, cast, get_args, get_origin, overload @@ -117,6 +119,28 @@ def _get_paginated_types( # --------------------------------------------------------------------------- +def _retry_after(response: httpx.Response) -> float | None: + """Parse a reasonable server-requested retry delay in seconds.""" + retry_after_ms = response.headers.get("retry-after-ms") + try: + delay = float(retry_after_ms) / 1000 + except (TypeError, ValueError): + retry_after = response.headers.get("retry-after") + try: + delay = float(retry_after) + except (TypeError, ValueError): + # Retry-After may instead be an RFC 5322 / HTTP-date, e.g. + # "Fri, 31 Dec 2027 23:59:59 GMT"; convert that to a delta. + try: + retry_date = email.utils.parsedate_to_datetime(retry_after) + except (TypeError, ValueError): + return None + if retry_date.tzinfo is None: + retry_date = retry_date.replace(tzinfo=timezone.utc) + delay = retry_date.timestamp() - time.time() + return delay if 0 < delay <= 60 else None + + def _should_retry( response: httpx.Response | None, exc: httpx.TransportError | None, @@ -129,14 +153,33 @@ def _should_retry( Returns the sleep duration if a retry should happen, or ``None`` if the response should be returned / the exception re-raised. """ - is_last = attempt >= policy.max_retries - if is_last: + if attempt >= policy.max_retries: return None + + backoff = policy.backoff_base * (2**attempt) if exc is not None: - return policy.backoff_base * (2**attempt) - if response is not None and response.status_code in policy.retryable_status_codes: - return policy.backoff_base * (2**attempt) - return None + return backoff + if response is None: + return None + + if policy.respect_retry_decision_headers: + if response.status_code < 400: + return None + should_retry = response.headers.get("x-should-retry") + if should_retry == "true": + return (_retry_after(response) or backoff) if policy.respect_retry_after_headers else backoff + if should_retry == "false": + return None + + retryable_status = response.status_code in policy.retryable_status_codes + if policy.retry_all_server_errors and response.status_code >= 500: + retryable_status = True + if not retryable_status: + return None + + if policy.respect_retry_after_headers: + return _retry_after(response) or backoff + return backoff def _should_resolve_conflict(response: httpx.Response, request: PreparedRequest) -> bool: diff --git a/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/types.py b/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/types.py index c53990986b..ad9667fb9d 100644 --- a/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/types.py +++ b/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/types.py @@ -277,8 +277,10 @@ class RetryPolicy: Set as a client-level default via the ``retry`` constructor parameter, or override per-request via ``send()``'s ``retry`` keyword argument. - This is an operational concern — it does not belong in endpoint - signatures. + This is an operational concern and does not belong in endpoint + signatures. Response-header handling and broad server-error retries are + opt-in so adapters can reproduce another client's retry contract without + changing standalone client defaults. .. note:: @@ -300,6 +302,9 @@ class RetryPolicy: max_retries: int = 3 backoff_base: float = 0.5 retryable_status_codes: tuple[int, ...] = (502, 503, 504, 429) + retry_all_server_errors: bool = False + respect_retry_decision_headers: bool = False + respect_retry_after_headers: bool = False @dataclass(frozen=True, slots=True) diff --git a/packages/nemo_platform_plugin/tests/client/test_adapter.py b/packages/nemo_platform_plugin/tests/client/test_adapter.py index e51c4fd19d..7d3f3d9e6f 100644 --- a/packages/nemo_platform_plugin/tests/client/test_adapter.py +++ b/packages/nemo_platform_plugin/tests/client/test_adapter.py @@ -6,11 +6,12 @@ import httpx from nemo_platform import NeMoPlatform from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.client.types import RetryPolicy from nemo_platform_plugin.jobs import endpoints from nemo_platform_plugin.jobs.client import JobsClient -def test_client_from_platform_preserves_retry_count_with_nemoclient_defaults() -> None: +def test_client_from_platform_preserves_stainless_retry_policy() -> None: http_client = httpx.Client(transport=httpx.MockTransport(lambda request: httpx.Response(200, request=request))) platform = NeMoPlatform( base_url="http://test", @@ -22,8 +23,13 @@ def test_client_from_platform_preserves_retry_count_with_nemoclient_defaults() - client = client_from_platform(platform, JobsClient) assert client.retry is not None - assert client.retry.max_retries == 4 - assert client.retry.retryable_status_codes == (502, 503, 504, 429) + assert client.retry == RetryPolicy( + max_retries=4, + retryable_status_codes=(408, 409, 429), + retry_all_server_errors=True, + respect_retry_decision_headers=True, + respect_retry_after_headers=True, + ) def test_client_from_platform_prefers_platform_request_router() -> None: diff --git a/packages/nemo_platform_plugin/tests/client/test_client_options.py b/packages/nemo_platform_plugin/tests/client/test_client_options.py index b0dee0fdf9..1dc4b83e94 100644 --- a/packages/nemo_platform_plugin/tests/client/test_client_options.py +++ b/packages/nemo_platform_plugin/tests/client/test_client_options.py @@ -5,7 +5,7 @@ from __future__ import annotations -from unittest.mock import AsyncMock, MagicMock +from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest @@ -23,6 +23,14 @@ from pydantic import BaseModel BASE = "http://test:8000" +STAINLESS_RETRY = RetryPolicy( + max_retries=1, + backoff_base=0.25, + retryable_status_codes=(408, 409, 429), + retry_all_server_errors=True, + respect_retry_decision_headers=True, + respect_retry_after_headers=True, +) class ItemRequest(BaseModel): @@ -459,6 +467,180 @@ def test_no_retry_without_policy(self) -> None: assert exc_info.value.status_code == 503 assert mock_http.request.call_count == 1 + def test_standalone_retry_policy_defaults_are_unchanged(self) -> None: + policy = RetryPolicy() + + assert policy.retryable_status_codes == (502, 503, 504, 429) + assert policy.retry_all_server_errors is False + assert policy.respect_retry_decision_headers is False + assert policy.respect_retry_after_headers is False + + @pytest.mark.parametrize("status_code", [408, 409, 500]) + def test_standalone_policy_does_not_add_stainless_statuses(self, status_code: int) -> None: + mock_http = MagicMock(spec=httpx.Client) + mock_http.request.return_value = httpx.Response( + status_code, + request=httpx.Request("GET", f"{BASE}/apis/test/v2/items/alice"), + json={"detail": "error"}, + ) + client = NemoClient( + base_url=BASE, + http_client=mock_http, + retry=RetryPolicy(max_retries=1, backoff_base=0.0), + ) + + with pytest.raises(NemoHTTPError): + client.send(GET_ITEM(name="alice")) + + assert mock_http.request.call_count == 1 + + @pytest.mark.parametrize("status_code", [408, 409, 500]) + def test_stainless_policy_retries_all_expected_statuses(self, status_code: int) -> None: + mock_http = MagicMock(spec=httpx.Client) + mock_http.request.side_effect = [ + httpx.Response( + status_code, + request=httpx.Request("GET", f"{BASE}/apis/test/v2/items/alice"), + json={"detail": "error"}, + ), + httpx.Response( + 200, + request=httpx.Request("GET", f"{BASE}/apis/test/v2/items/alice"), + json={"id": 1, "name": "alice"}, + ), + ] + client = NemoClient(base_url=BASE, http_client=mock_http, retry=STAINLESS_RETRY) + + with patch("nemo_platform_plugin.client.client.time.sleep"): + response = client.send(GET_ITEM(name="alice")) + + assert response.body.name == "alice" + assert mock_http.request.call_count == 2 + + def test_stainless_true_header_forces_retry(self) -> None: + mock_http = MagicMock(spec=httpx.Client) + mock_http.request.side_effect = [ + httpx.Response( + 400, + headers={"x-should-retry": "true"}, + request=httpx.Request("GET", f"{BASE}/apis/test/v2/items/alice"), + json={"detail": "retry"}, + ), + httpx.Response( + 200, + request=httpx.Request("GET", f"{BASE}/apis/test/v2/items/alice"), + json={"id": 1, "name": "alice"}, + ), + ] + client = NemoClient(base_url=BASE, http_client=mock_http, retry=STAINLESS_RETRY) + + with patch("nemo_platform_plugin.client.client.time.sleep"): + response = client.send(GET_ITEM(name="alice")) + + assert response.body.name == "alice" + assert mock_http.request.call_count == 2 + + def test_stainless_true_header_does_not_retry_success(self) -> None: + mock_http = MagicMock(spec=httpx.Client) + mock_http.request.return_value = httpx.Response( + 200, + headers={"x-should-retry": "true"}, + request=httpx.Request("GET", f"{BASE}/apis/test/v2/items/alice"), + json={"id": 1, "name": "alice"}, + ) + client = NemoClient(base_url=BASE, http_client=mock_http, retry=STAINLESS_RETRY) + + response = client.send(GET_ITEM(name="alice")) + + assert response.body.name == "alice" + assert mock_http.request.call_count == 1 + + def test_stainless_false_header_suppresses_retry(self) -> None: + mock_http = MagicMock(spec=httpx.Client) + mock_http.request.return_value = httpx.Response( + 500, + headers={"x-should-retry": "false"}, + request=httpx.Request("GET", f"{BASE}/apis/test/v2/items/alice"), + json={"detail": "do not retry"}, + ) + client = NemoClient(base_url=BASE, http_client=mock_http, retry=STAINLESS_RETRY) + + with pytest.raises(NemoHTTPError): + client.send(GET_ITEM(name="alice")) + + assert mock_http.request.call_count == 1 + + @pytest.mark.parametrize( + ("headers", "expected_delay"), + [ + ({"retry-after-ms": "1250"}, 1.25), + ({"retry-after-ms": "invalid", "retry-after": "3"}, 3.0), + ({"retry-after": "2.5"}, 2.5), + ({"retry-after": "Mon, 12 Jan 1970 13:47:10 GMT"}, 30.0), + ({"retry-after": "60"}, 60.0), + ], + ) + def test_stainless_policy_honors_retry_after(self, headers: dict[str, str], expected_delay: float) -> None: + mock_http = MagicMock(spec=httpx.Client) + mock_http.request.side_effect = [ + httpx.Response( + 500, + headers=headers, + request=httpx.Request("GET", f"{BASE}/apis/test/v2/items/alice"), + json={"detail": "retry"}, + ), + httpx.Response( + 200, + request=httpx.Request("GET", f"{BASE}/apis/test/v2/items/alice"), + json={"id": 1, "name": "alice"}, + ), + ] + client = NemoClient(base_url=BASE, http_client=mock_http, retry=STAINLESS_RETRY) + + with ( + patch("nemo_platform_plugin.client.client.time.time", return_value=1_000_000), + patch("nemo_platform_plugin.client.client.time.sleep") as sleep, + ): + client.send(GET_ITEM(name="alice")) + + sleep.assert_called_once_with(expected_delay) + + @pytest.mark.parametrize( + "headers", + [ + {"retry-after-ms": "0"}, + {"retry-after-ms": "60001"}, + {"retry-after": "-1"}, + {"retry-after": "60.1"}, + {"retry-after": "not-a-delay"}, + {"retry-after": "Mon, 12 Jan 1970 13:46:39 GMT"}, + ], + ) + def test_stainless_policy_falls_back_for_unreasonable_retry_after(self, headers: dict[str, str]) -> None: + mock_http = MagicMock(spec=httpx.Client) + mock_http.request.side_effect = [ + httpx.Response( + 500, + headers=headers, + request=httpx.Request("GET", f"{BASE}/apis/test/v2/items/alice"), + json={"detail": "retry"}, + ), + httpx.Response( + 200, + request=httpx.Request("GET", f"{BASE}/apis/test/v2/items/alice"), + json={"id": 1, "name": "alice"}, + ), + ] + client = NemoClient(base_url=BASE, http_client=mock_http, retry=STAINLESS_RETRY) + + with ( + patch("nemo_platform_plugin.client.client.time.time", return_value=1_000_000), + patch("nemo_platform_plugin.client.client.time.sleep") as sleep, + ): + client.send(GET_ITEM(name="alice")) + + sleep.assert_called_once_with(STAINLESS_RETRY.backoff_base) + def test_binary_stream_retries_before_returning_content(self) -> None: attempts = 0 @@ -540,6 +722,86 @@ async def test_retry_on_503_async(self) -> None: assert resp.body.name == "alice" assert mock_http.request.call_count == 2 + @pytest.mark.asyncio + @pytest.mark.parametrize("status_code", [408, 409, 500]) + async def test_stainless_policy_retries_expected_statuses_async(self, status_code: int) -> None: + mock_http = AsyncMock(spec=httpx.AsyncClient) + mock_http.request.side_effect = [ + httpx.Response( + status_code, + request=httpx.Request("GET", f"{BASE}/apis/test/v2/items/alice"), + json={"detail": "retry"}, + ), + httpx.Response( + 200, + request=httpx.Request("GET", f"{BASE}/apis/test/v2/items/alice"), + json={"id": 1, "name": "alice"}, + ), + ] + client = AsyncNemoClient(base_url=BASE, http_client=mock_http, retry=STAINLESS_RETRY) + + with patch("nemo_platform_plugin.client.client.asyncio.sleep", new_callable=AsyncMock): + response = await client.send(GET_ITEM(name="alice")) + + assert response.body.name == "alice" + assert mock_http.request.call_count == 2 + + @pytest.mark.asyncio + async def test_stainless_true_header_forces_retry_async(self) -> None: + mock_http = AsyncMock(spec=httpx.AsyncClient) + mock_http.request.side_effect = [ + httpx.Response( + 400, + headers={"x-should-retry": "true"}, + request=httpx.Request("GET", f"{BASE}/apis/test/v2/items/alice"), + json={"detail": "retry"}, + ), + httpx.Response( + 200, + request=httpx.Request("GET", f"{BASE}/apis/test/v2/items/alice"), + json={"id": 1, "name": "alice"}, + ), + ] + client = AsyncNemoClient(base_url=BASE, http_client=mock_http, retry=STAINLESS_RETRY) + + with patch("nemo_platform_plugin.client.client.asyncio.sleep", new_callable=AsyncMock): + response = await client.send(GET_ITEM(name="alice")) + + assert response.body.name == "alice" + assert mock_http.request.call_count == 2 + + @pytest.mark.asyncio + async def test_stainless_true_header_does_not_retry_success_async(self) -> None: + mock_http = AsyncMock(spec=httpx.AsyncClient) + mock_http.request.return_value = httpx.Response( + 200, + headers={"x-should-retry": "true"}, + request=httpx.Request("GET", f"{BASE}/apis/test/v2/items/alice"), + json={"id": 1, "name": "alice"}, + ) + client = AsyncNemoClient(base_url=BASE, http_client=mock_http, retry=STAINLESS_RETRY) + + response = await client.send(GET_ITEM(name="alice")) + + assert response.body.name == "alice" + assert mock_http.request.call_count == 1 + + @pytest.mark.asyncio + async def test_stainless_false_header_suppresses_retry_async(self) -> None: + mock_http = AsyncMock(spec=httpx.AsyncClient) + mock_http.request.return_value = httpx.Response( + 500, + headers={"x-should-retry": "false"}, + request=httpx.Request("GET", f"{BASE}/apis/test/v2/items/alice"), + json={"detail": "do not retry"}, + ) + client = AsyncNemoClient(base_url=BASE, http_client=mock_http, retry=STAINLESS_RETRY) + + with pytest.raises(NemoHTTPError): + await client.send(GET_ITEM(name="alice")) + + assert mock_http.request.call_count == 1 + @pytest.mark.asyncio async def test_exhausted_transport_error_is_wrapped_async(self) -> None: mock_http = AsyncMock(spec=httpx.AsyncClient)