Skip to content

fix(client): preserve SDK retry behavior in NemoClient adapters #992

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
maxdubrinsky wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix(client): preserve SDK retry behavior in NemoClient adapters #992

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 7 additions & 1 deletion packages/nemo_platform_plugin/src/nemo_platform_plugin/client/adapter.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down
55 changes: 49 additions & 6 deletions packages/nemo_platform_plugin/src/nemo_platform_plugin/client/client.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Comment on lines +134 to +141

@coderabbitai coderabbitai Bot Jul 30, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win

Guard against OverflowError/OSError from far-future HTTP-dates.

parsedate_to_datetime and datetime.timestamp() can raise OverflowError (or OSError on some platforms) for out-of-range dates, escaping _retry_after and aborting the retry loop with an unrelated exception.

🛡️ Proposed fix 
             try:
                 retry_date = email.utils.parsedate_to_datetime(retry_after)
-            except (TypeError, ValueError):
+            except (TypeError, ValueError, OverflowError):
                 return None
             if retry_date.tzinfo is None:
                 retry_date = retry_date.replace(tzinfo=timezone.utc)
-            delay = retry_date.timestamp() - time.time()
+            try:
+                delay = retry_date.timestamp() - time.time()
+            except (OverflowError, OSError, ValueError):
+                return None
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
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
try:
retry_date = email.utils.parsedate_to_datetime(retry_after)
except (TypeError, ValueError, OverflowError):
return None
if retry_date.tzinfo is None:
retry_date = retry_date.replace(tzinfo=timezone.utc)
try:
delay = retry_date.timestamp() - time.time()
except (OverflowError, OSError, ValueError):
return None
return delay if 0 < delay <= 60 else None
🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@packages/nemo_platform_plugin/src/nemo_platform_plugin/client/client.py`
around lines 134 - 141, Update _retry_after to catch OverflowError and OSError
alongside TypeError and ValueError when parsing retry_after and converting the
resulting datetime to a timestamp. Return None for these out-of-range HTTP-date
cases so the retry loop continues without propagating the exception.



def _should_retry(
response: httpx.Response | None,
exc: httpx.TransportError | None,
Expand All @@ -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:
Expand Down
9 changes: 7 additions & 2 deletions packages/nemo_platform_plugin/src/nemo_platform_plugin/client/types.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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::

Expand All @@ -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)
Expand Down
12 changes: 9 additions & 3 deletions packages/nemo_platform_plugin/tests/client/test_adapter.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand All @@ -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:
Expand Down
Loading
Loading