diff --git a/packages/filesets/src/filesets/filesystem/filesystem.py b/packages/filesets/src/filesets/filesystem/filesystem.py index 9862c0def3..ac400f7dc9 100644 --- a/packages/filesets/src/filesets/filesystem/filesystem.py +++ b/packages/filesets/src/filesets/filesystem/filesystem.py @@ -358,24 +358,24 @@ def _ensure_async(client: FilesClient | AsyncFilesClient) -> AsyncFilesClient: import httpx + # A timeout lives in two layers, so mirror each from its own source: the + # transport carries the client-level default, and ``_timeout`` the + # per-request override that ``send`` puts on every request. Leave the + # transport's unset and httpx falls back to its own 5s, which a multi-GB + # upload blows through waiting for the server to commit the body to storage. asgi_app = getattr(client._http, "asgi_app", None) - http_client = ( - httpx.AsyncClient( - transport=httpx.ASGITransport(app=asgi_app), - base_url=client.base_url, - headers=dict(client._default_headers) if client._default_headers else None, - ) - if asgi_app is not None - else httpx.AsyncClient( - base_url=client.base_url, - headers=dict(client._default_headers) if client._default_headers else None, - ) + http_client = httpx.AsyncClient( + transport=httpx.ASGITransport(app=asgi_app) if asgi_app is not None else None, + base_url=client.base_url, + headers=dict(client._default_headers) if client._default_headers else None, + timeout=client._http.timeout, ) return AsyncFilesClient( base_url=client.base_url, workspace=client.workspace, auth=client._auth, default_headers=client._default_headers or None, + timeout=client._timeout, retry=client._retry, http_client=http_client, url_resolver=client._url_resolver, diff --git a/packages/filesets/tests/test_filesystem_client.py b/packages/filesets/tests/test_filesystem_client.py new file mode 100644 index 0000000000..2b6ad2d894 --- /dev/null +++ b/packages/filesets/tests/test_filesystem_client.py @@ -0,0 +1,80 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Client construction inside FilesetFileSystem. + +Uploads and downloads run on the async client built by ``_ensure_async``, not on +the sync client the caller configured. Anything that client fails to carry over +is silently dropped from every transfer. +""" + +from __future__ import annotations + +import httpx +from filesets.filesystem.filesystem import FilesetFileSystem +from nemo_platform_plugin.client.types import RetryPolicy +from nemo_platform_plugin.files.client import FilesClient + +BASE = "http://test:8000" +UPLOAD_TIMEOUT = httpx.Timeout(30.0, write=10 * 60, read=5 * 60) + + +def _sync_client(*, timeout: httpx.Timeout) -> FilesClient: + http_client = httpx.Client( + transport=httpx.MockTransport(lambda request: httpx.Response(200, request=request)), + timeout=timeout, + ) + return FilesClient( + base_url=BASE, + workspace="default", + http_client=http_client, + retry=RetryPolicy(max_retries=2), + ) + + +def test_ensure_async_carries_transport_timeout() -> None: + """With no override to carry, the transport's own timeout is what governs. + + Without this the new AsyncClient falls back to httpx's 5s default. + """ + async_client = FilesetFileSystem._ensure_async(_sync_client(timeout=httpx.Timeout(60.0))) + + assert async_client._timeout is None + assert async_client._http.timeout == httpx.Timeout(60.0) + assert async_client._http.timeout != httpx.Timeout(5.0) + + +def test_ensure_async_carries_per_request_timeout_override() -> None: + """An override goes out on every request, so it governs regardless of the transport.""" + client = _sync_client(timeout=httpx.Timeout(60.0)).with_options(timeout=UPLOAD_TIMEOUT) + + async_client = FilesetFileSystem._ensure_async(client) + + assert async_client._timeout == UPLOAD_TIMEOUT + # Each layer is copied from its own counterpart, so the transport keeps the + # client-level default it had on the sync side rather than the override. + assert async_client._http.timeout == httpx.Timeout(60.0) + + +def test_ensure_async_preserves_workspace_and_retry() -> None: + client = _sync_client(timeout=httpx.Timeout(60.0)) + + async_client = FilesetFileSystem._ensure_async(client) + + assert async_client.workspace == "default" + assert async_client.retry == RetryPolicy(max_retries=2) + + +def test_upload_timeout_survives_the_whole_client_chain() -> None: + """End to end: an SDK-level timeout override reaches the client that transfers.""" + from nemo_platform import NeMoPlatform + + http_client = httpx.Client( + transport=httpx.MockTransport(lambda request: httpx.Response(200, request=request)), + timeout=httpx.Timeout(60.0), + ) + platform = NeMoPlatform(base_url=BASE, workspace="default", http_client=http_client) + + fs = platform.with_options(timeout=UPLOAD_TIMEOUT).files.fsspec + + assert fs._client._timeout == UPLOAD_TIMEOUT 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 dd177bd7dc..410a1ae6fc 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 @@ -69,6 +69,19 @@ def client_from_platform( respect_retry_after_headers=True, ) url_resolver = _url_resolver_from_platform(platform) + + # Carry the platform's timeout across as a per-request override. The shared + # httpx client keeps whatever timeout it was built with, so a caller's + # ``platform.with_options(timeout=...)`` would otherwise be silently dropped + # on the way to the typed client — the httpx client it hands over is the + # *same* object, with the *original* timeout still on it. + timeout = platform.timeout + if timeout is None: + # ``None`` on the platform means "no timeout at all", but the typed + # client reads None as "defer to the transport". Say the same thing in + # the form httpx itself uses, so the override survives. + timeout = httpx.Timeout(None) + if isinstance(platform, AsyncNeMoPlatform): if not issubclass(client_cls, AsyncNemoClient): raise TypeError("AsyncNeMoPlatform requires an AsyncNemoClient class") @@ -76,16 +89,19 @@ def client_from_platform( base_url=str(platform.base_url).rstrip("/"), workspace=platform.workspace, default_headers=headers or None, + timeout=timeout, retry=retry, http_client=platform._client, url_resolver=url_resolver, ) + if not issubclass(client_cls, NemoClient): raise TypeError("NeMoPlatform requires a NemoClient class") return client_cls( base_url=str(platform.base_url).rstrip("/"), workspace=platform.workspace, default_headers=headers or None, + timeout=timeout, retry=retry, http_client=platform._client, url_resolver=url_resolver, 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 3fb6add084..65e994b9c4 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 @@ -21,9 +21,10 @@ import email.utils import inspect import json +import logging import os import time -from collections.abc import AsyncIterator, Callable, Iterator, Mapping +from collections.abc import AsyncIterable, AsyncIterator, Callable, Iterable, Iterator, Mapping from contextlib import asynccontextmanager, contextmanager from datetime import timezone from functools import cache @@ -68,6 +69,8 @@ ModelT = TypeVar("ModelT", bound=BaseModel) +logger = logging.getLogger(__name__) + DEFAULT_TIMEOUT = 60.0 @@ -144,7 +147,45 @@ def _retry_after(response: httpx.Response) -> float | None: return delay if 0 < delay <= 60 else None +# Transport failures httpx raises before it starts reading the request body: the +# connection has to exist before any body byte can be written. A one-shot body is +# therefore still untouched when one of these surfaces, so the request can be sent +# again as-is. Everything else — a read/write timeout, a dropped connection, a +# protocol error — can land mid-body, where a replay would send a truncated body +# under the original ``Content-Length``. +_PRE_BODY_TRANSPORT_ERRORS = ( + httpx.ConnectError, + httpx.ConnectTimeout, + httpx.PoolTimeout, + httpx.ProxyError, + httpx.UnsupportedProtocol, +) + + +def _is_replayable(content: bytes | Iterable[bytes] | AsyncIterable[bytes] | None) -> bool: + """Whether a request body can be handed to httpx more than once. + + ``bytes``, ``str``, no body at all, and in-memory sequences replay fine — + httpx re-reads them from the start on every attempt. + + A generator, file object or other one-shot iterable does not: httpx drains it + on the first attempt, so a replay sends a short body while the original + ``Content-Length`` still stands on the request, and h11 aborts with ``Too + little data for declared Content-Length`` — masking whatever actually failed + the first time. Such a body may only be re-sent while it is still untouched; + once it has been read, the retry has to happen a level up, where the caller + can build a fresh iterator over the source. + + ``bytearray`` and ``memoryview`` are deliberately absent: httpx treats + anything that is not ``bytes``/``str`` as an iterable of chunks, and + iterating either of those yields ``int``, so it rejects them on the first + attempt regardless of what this returns. + """ + return content is None or isinstance(content, (bytes, str, list, tuple)) + + def _should_retry( + request: PreparedRequest, response: httpx.Response | None, exc: httpx.TransportError | None, attempt: int, @@ -155,33 +196,49 @@ def _should_retry( Shared decision logic used by both sync and async retry paths. Returns the sleep duration if a retry should happen, or ``None`` if the response should be returned / the exception re-raised. + + The policy decides first, without regard for the body; only then does the + body get a veto. Keeping that order means a one-shot body (see + :func:`_is_replayable`) is reported as the reason a retry stopped exactly + when it is the reason, rather than on every attempt that was never going to + be retried anyway. """ if attempt >= policy.max_retries: return None backoff = policy.backoff_base * (2**attempt) if exc is not None: - return backoff - if response is None: + # A connection that never opened leaves the body untouched; anything + # later can land mid-body. See :data:`_PRE_BODY_TRANSPORT_ERRORS`. + body_is_spent = not isinstance(exc, _PRE_BODY_TRANSPORT_ERRORS) + elif response is None: return None - - if policy.respect_retry_decision_headers: - if response.status_code < 400: + else: + # A response only arrives once the body has gone out on the wire. + body_is_spent = True + decision = response.headers.get("x-should-retry") if policy.respect_retry_decision_headers else None + if policy.respect_retry_decision_headers and 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": + if decision == "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: + if decision != "true": + # No explicit verdict from the server, so fall back to the status code. + 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: + backoff = _retry_after(response) or backoff + + if body_is_spent and not _is_replayable(request.content): + logger.info( + "Not retrying %s %s: the request body is a one-shot stream that has already been read, " + "so it cannot be sent again. Retry at a level that can rebuild it.", + request.method, + request.path_template, + ) return None - - if policy.respect_retry_after_headers: - return _retry_after(response) or backoff return backoff @@ -221,6 +278,7 @@ def __init__( auth: TokenProvider | AsyncTokenProvider | str | None = None, retry: RetryPolicy | None = None, default_headers: Mapping[str, str] | None = None, + timeout: float | httpx.Timeout | None = None, url_resolver: Callable[[str], str | httpx.URL] | None = None, ) -> None: self._base_url = base_url.rstrip("/") @@ -229,7 +287,7 @@ def __init__( self._retry = retry self._default_headers = dict(default_headers) if default_headers else {} self._url_resolver = url_resolver - self._timeout: float | httpx.Timeout | None = None + self._timeout: float | httpx.Timeout | None = timeout @property def base_url(self) -> str: @@ -349,22 +407,31 @@ def __init__( workspace: str | None = None, auth: TokenProvider | str | None = None, default_headers: Mapping[str, str] | None = None, - timeout: float = DEFAULT_TIMEOUT, + timeout: float | httpx.Timeout | None = None, retry: RetryPolicy | None = None, http_client: httpx.Client | None = None, url_resolver: Callable[[str], str | httpx.URL] | None = None, ) -> None: + """Create a client. + + *timeout* is applied per request, so it holds even when *http_client* is + supplied and shared with another caller — an httpx client's own timeout + is fixed when it is built and cannot be changed afterwards. ``None`` + defers to the transport's timeout, giving one we build ourselves + :data:`DEFAULT_TIMEOUT`; ``httpx.Timeout(None)`` waits indefinitely. + """ super().__init__( base_url=base_url, workspace=workspace, auth=auth, retry=retry, default_headers=default_headers, + timeout=timeout, url_resolver=url_resolver, ) self._http = http_client or httpx.Client( headers=dict(default_headers) if default_headers else None, - timeout=timeout, + timeout=timeout if timeout is not None else DEFAULT_TIMEOUT, ) @classmethod @@ -375,6 +442,7 @@ def from_client(cls, client: NemoClient) -> Self: workspace=client.workspace, auth=client._auth, default_headers=client._default_headers or None, + timeout=client._timeout, retry=client._retry, http_client=client._http, url_resolver=client._url_resolver, @@ -519,13 +587,13 @@ def _request_with_retry( kwargs["timeout"] = self._timeout raw = self._http.request(request.method, url, **kwargs) except httpx.TransportError as exc: - backoff = _should_retry(None, exc, attempt, retry) if retry else None + backoff = _should_retry(request, None, exc, attempt, retry) if retry else None if backoff is not None: time.sleep(backoff) continue raise NemoTransportError(exc) from exc if retry: - backoff = _should_retry(raw, None, attempt, retry) + backoff = _should_retry(request, raw, None, attempt, retry) if backoff is not None: last_response = raw time.sleep(backoff) @@ -552,7 +620,7 @@ def _stream_with_retry( if self._timeout is not None: kwargs["timeout"] = self._timeout with self._http.stream(request.method, url, **kwargs) as raw: - backoff = _should_retry(raw, None, attempt, retry) if retry else None + backoff = _should_retry(request, raw, None, attempt, retry) if retry else None if backoff is not None: time.sleep(backoff) continue @@ -562,7 +630,7 @@ def _stream_with_retry( except httpx.TransportError as exc: if yielded: raise NemoTransportError(exc) from exc - backoff = _should_retry(None, exc, attempt, retry) if retry else None + backoff = _should_retry(request, None, exc, attempt, retry) if retry else None if backoff is not None: time.sleep(backoff) continue @@ -597,22 +665,24 @@ def __init__( workspace: str | None = None, auth: TokenProvider | AsyncTokenProvider | str | None = None, default_headers: Mapping[str, str] | None = None, - timeout: float = DEFAULT_TIMEOUT, + timeout: float | httpx.Timeout | None = None, retry: RetryPolicy | None = None, http_client: httpx.AsyncClient | None = None, url_resolver: Callable[[str], str | httpx.URL] | None = None, ) -> None: + """Create a client. See :meth:`NemoClient.__init__` for *timeout*.""" super().__init__( base_url=base_url, workspace=workspace, auth=auth, retry=retry, default_headers=default_headers, + timeout=timeout, url_resolver=url_resolver, ) self._http = http_client or httpx.AsyncClient( headers=dict(default_headers) if default_headers else None, - timeout=timeout, + timeout=timeout if timeout is not None else DEFAULT_TIMEOUT, ) @classmethod @@ -623,6 +693,7 @@ def from_client(cls, client: AsyncNemoClient) -> Self: workspace=client.workspace, auth=client._auth, default_headers=client._default_headers or None, + timeout=client._timeout, retry=client._retry, http_client=client._http, url_resolver=client._url_resolver, @@ -759,13 +830,13 @@ async def _request_with_retry( kwargs["timeout"] = self._timeout raw = await self._http.request(request.method, url, **kwargs) except httpx.TransportError as exc: - backoff = _should_retry(None, exc, attempt, retry) if retry else None + backoff = _should_retry(request, None, exc, attempt, retry) if retry else None if backoff is not None: await asyncio.sleep(backoff) continue raise NemoTransportError(exc) from exc if retry: - backoff = _should_retry(raw, None, attempt, retry) + backoff = _should_retry(request, raw, None, attempt, retry) if backoff is not None: last_response = raw await asyncio.sleep(backoff) @@ -792,7 +863,7 @@ async def _stream_with_retry( if self._timeout is not None: kwargs["timeout"] = self._timeout async with self._http.stream(request.method, url, **kwargs) as raw: - backoff = _should_retry(raw, None, attempt, retry) if retry else None + backoff = _should_retry(request, raw, None, attempt, retry) if retry else None if backoff is not None: await asyncio.sleep(backoff) continue @@ -802,7 +873,7 @@ async def _stream_with_retry( except httpx.TransportError as exc: if yielded: raise NemoTransportError(exc) from exc - backoff = _should_retry(None, exc, attempt, retry) if retry else None + backoff = _should_retry(request, None, exc, attempt, retry) if retry else None if backoff is not None: await asyncio.sleep(backoff) continue 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 ad9667fb9d..e42f17ba54 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 @@ -290,6 +290,15 @@ class RetryPolicy: ``Idempotency-Key`` header for create operations that must be safe to retry. + A request whose ``content`` is a generator, file object or other + one-shot iterable is only retried while its body is still + untouched — that is, on a connection-establishment failure. Once + httpx has started reading the body it cannot be sent again, so a + mid-body failure (and any retryable status code, which by + definition arrives after the body) is raised to the caller + instead. Retry those where the body can be rebuilt from its + source. + Usage:: # Client-level default diff --git a/packages/nemo_platform_plugin/tests/client/test_adapter.py b/packages/nemo_platform_plugin/tests/client/test_adapter.py index 7d3f3d9e6f..7f82c0e985 100644 --- a/packages/nemo_platform_plugin/tests/client/test_adapter.py +++ b/packages/nemo_platform_plugin/tests/client/test_adapter.py @@ -44,7 +44,7 @@ class RequestRouter: def resolve(self, url: str) -> str: return url.replace("http://gateway/apis/jobs", "http://127.0.0.1:8080/apis/jobs") - platform._nmp_request_router = RequestRouter() # type: ignore[attr-defined] + platform._nmp_request_router = RequestRouter() # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] client = client_from_platform(platform, JobsClient) @@ -63,7 +63,7 @@ def test_client_from_platform_falls_back_to_sdk_prepare_url() -> None: def prepare_url(url: str) -> str: return url.replace("http://gateway/apis/jobs", "http://127.0.0.1:8080/apis/jobs") - platform._prepare_url = prepare_url # type: ignore[method-assign] + platform._prepare_url = prepare_url # type: ignore[method-assign] # ty: ignore[invalid-assignment] client = client_from_platform(platform, JobsClient) @@ -84,3 +84,51 @@ def test_from_client_preserves_url_resolver() -> None: request = endpoints.list_steps(workspace="default", name="job-1") assert clone._resolve_path(request) == ("http://127.0.0.1:8080/apis/jobs/v2/workspaces/default/jobs/job-1/steps") + + +def test_client_from_platform_propagates_timeout() -> None: + """``platform.with_options(timeout=...)`` must reach the typed client. + + Both clients share one httpx client, whose own timeout ``with_options`` does + not touch — so the typed client has to carry the override itself or long + transfers silently run on the transport's original budget. + """ + upload_timeout = httpx.Timeout(30.0, write=10 * 60, read=5 * 60) + http_client = httpx.Client( + transport=httpx.MockTransport(lambda request: httpx.Response(200, request=request)), + timeout=httpx.Timeout(60.0), + ) + platform = NeMoPlatform(base_url="http://test", workspace="default", http_client=http_client) + + scoped = platform.with_options(timeout=upload_timeout) + client = client_from_platform(scoped, JobsClient) + + assert client._timeout == upload_timeout + # The shared transport is untouched, which is why the override is needed. + assert scoped._client.timeout == httpx.Timeout(60.0) + + +def test_client_from_platform_carries_default_timeout() -> None: + http_client = httpx.Client( + transport=httpx.MockTransport(lambda request: httpx.Response(200, request=request)), + timeout=httpx.Timeout(60.0), + ) + platform = NeMoPlatform(base_url="http://test", workspace="default", http_client=http_client) + + client = client_from_platform(platform, JobsClient) + + assert client._timeout == platform.timeout + + +def test_client_from_platform_carries_disabled_timeout() -> None: + """``timeout=None`` means "no timeout", not "no override".""" + http_client = httpx.Client( + transport=httpx.MockTransport(lambda request: httpx.Response(200, request=request)), + timeout=httpx.Timeout(60.0), + ) + platform = NeMoPlatform(base_url="http://test", workspace="default", http_client=http_client) + + client = client_from_platform(platform.with_options(timeout=None), JobsClient) + + # Not the transport's 60s: httpx reads an all-None Timeout as "wait forever". + assert client._timeout == httpx.Timeout(None) diff --git a/packages/nemo_platform_plugin/tests/client/test_client.py b/packages/nemo_platform_plugin/tests/client/test_client.py index cbd1ee079e..1e76d9f616 100644 --- a/packages/nemo_platform_plugin/tests/client/test_client.py +++ b/packages/nemo_platform_plugin/tests/client/test_client.py @@ -7,7 +7,7 @@ import httpx import pytest -from nemo_platform_plugin.client.client import AsyncNemoClient, NemoClient, _type_adapter +from nemo_platform_plugin.client.client import DEFAULT_TIMEOUT, AsyncNemoClient, NemoClient, _type_adapter from nemo_platform_plugin.client.endpoint import delete, get, post from nemo_platform_plugin.client.errors import NemoHTTPError, NemoResponseValidationError, NotFoundError from nemo_platform_plugin.client.response import NemoResponse @@ -470,6 +470,75 @@ def test_response_carries_prepared_request() -> None: assert resp.request.path_params == {"name": "alice"} +# --------------------------------------------------------------------------- +# Timeout +# --------------------------------------------------------------------------- + + +def test_constructor_timeout_is_sent_with_every_request() -> None: + """A supplied transport's timeout is fixed, so the override rides per request.""" + mock_http = MagicMock(spec=httpx.Client) + mock_http.request.return_value = httpx.Response( + 200, + request=httpx.Request("GET", f"{BASE}/apis/test/v2/items/alice"), + json={"id": 1, "name": "alice"}, + ) + upload_timeout = httpx.Timeout(30.0, write=10 * 60, read=5 * 60) + + client = NemoClient(base_url=BASE, http_client=mock_http, timeout=upload_timeout) + client.send(GET_ITEM(name="alice")) + + assert mock_http.request.call_args.kwargs["timeout"] == upload_timeout + + +def test_omitted_timeout_defers_to_the_transport() -> None: + mock_http = MagicMock(spec=httpx.Client) + mock_http.request.return_value = 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) + client.send(GET_ITEM(name="alice")) + + assert "timeout" not in mock_http.request.call_args.kwargs + + +def test_owned_transport_is_built_with_the_default_timeout() -> None: + client = NemoClient(base_url=BASE) + + assert client._http.timeout == httpx.Timeout(DEFAULT_TIMEOUT) + + +def test_from_client_carries_the_timeout() -> None: + """The clone shares the transport, so it must carry the override too.""" + upload_timeout = httpx.Timeout(30.0, write=10 * 60, read=5 * 60) + client = NemoClient( + base_url=BASE, + http_client=MagicMock(spec=httpx.Client), + timeout=upload_timeout, + ) + + assert NemoClient.from_client(client)._timeout == upload_timeout + + +@pytest.mark.asyncio +async def test_constructor_timeout_is_sent_with_every_request_async() -> None: + mock_http = AsyncMock(spec=httpx.AsyncClient) + mock_http.request.return_value = httpx.Response( + 200, + request=httpx.Request("GET", f"{BASE}/apis/test/v2/items/alice"), + json={"id": 1, "name": "alice"}, + ) + upload_timeout = httpx.Timeout(30.0, write=10 * 60, read=5 * 60) + + client = AsyncNemoClient(base_url=BASE, http_client=mock_http, timeout=upload_timeout) + await client.send(GET_ITEM(name="alice")) + + assert mock_http.request.call_args.kwargs["timeout"] == upload_timeout + + # --------------------------------------------------------------------------- # Regression tests # --------------------------------------------------------------------------- diff --git a/packages/nemo_platform_plugin/tests/client/test_retry_streaming_body.py b/packages/nemo_platform_plugin/tests/client/test_retry_streaming_body.py new file mode 100644 index 0000000000..f42f6820fc --- /dev/null +++ b/packages/nemo_platform_plugin/tests/client/test_retry_streaming_body.py @@ -0,0 +1,427 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Retry behaviour for requests whose body cannot be replayed. + +A streaming upload sends a one-shot iterator under an explicit ``Content-Length``. +Replaying that request re-sends an exhausted iterator under the original length, +which h11 rejects with ``Too little data for declared Content-Length`` — hiding +whatever actually failed on the first attempt. + +That only applies once httpx has begun reading the body. A failure to establish +the connection leaves the iterator untouched, so those attempts are still retried. +""" + +from __future__ import annotations + +import asyncio +import logging +from collections.abc import AsyncIterable, AsyncIterator, Iterable +from unittest.mock import AsyncMock, MagicMock + +import httpx +import pytest +from nemo_platform_plugin.client.client import AsyncNemoClient, NemoClient +from nemo_platform_plugin.client.errors import InternalServerError, NemoTransportError, NotFoundError +from nemo_platform_plugin.client.types import PreparedRequest, RetryPolicy + +BASE = "http://test:8000" +PAYLOAD = 256 * 1024 +# Long enough to outlast the client's read timeout; the handler task is cancelled +# on teardown so the test never actually waits this out. +STALL_SECONDS = 30 +CLIENT_LOGGER = "nemo_platform_plugin.client.client" +# What ``client_from_platform`` builds for every typed client in production, and +# the only policy under which the response-header branches below are reachable. +ADAPTER_POLICY = RetryPolicy( + max_retries=3, + backoff_base=0.0, + retryable_status_codes=(408, 409, 429), + retry_all_server_errors=True, + respect_retry_decision_headers=True, + respect_retry_after_headers=True, +) + + +def _upload_request(content: bytes | Iterable[bytes] | AsyncIterable[bytes] | None) -> PreparedRequest: + return PreparedRequest( + path_template="/apis/test/v2/upload", + path_params={}, + method="PUT", + content=content, + content_type="application/octet-stream", + response_type=None, + ) + + +async def _chunks() -> AsyncIterator[bytes]: + sent = 0 + while sent < PAYLOAD: + n = min(64 * 1024, PAYLOAD - sent) + sent += n + yield b"x" * n + + +def _sync_chunks(): + sent = 0 + while sent < PAYLOAD: + n = min(64 * 1024, PAYLOAD - sent) + sent += n + yield b"x" * n + + +# --------------------------------------------------------------------------- +# The retry loop must not replay a body it has already started sending +# --------------------------------------------------------------------------- + + +async def test_async_streaming_body_is_not_replayed() -> None: + mock_http = MagicMock(spec=httpx.AsyncClient) + mock_http.request = AsyncMock(side_effect=httpx.ReadTimeout("timed out")) + + client = AsyncNemoClient( + base_url=BASE, + http_client=mock_http, + retry=RetryPolicy(max_retries=3, backoff_base=0.0), + ) + + with pytest.raises(NemoTransportError): + await client.with_headers({"Content-Length": str(PAYLOAD)}).send(_upload_request(_chunks())) + + assert mock_http.request.await_count == 1 + + +def test_sync_streaming_body_is_not_replayed() -> None: + mock_http = MagicMock(spec=httpx.Client) + mock_http.request.side_effect = httpx.ReadTimeout("timed out") + + client = NemoClient( + base_url=BASE, + http_client=mock_http, + retry=RetryPolicy(max_retries=3, backoff_base=0.0), + ) + + with pytest.raises(NemoTransportError): + client.with_headers({"Content-Length": str(PAYLOAD)}).send(_upload_request(_sync_chunks())) + + assert mock_http.request.call_count == 1 + + +# --------------------------------------------------------------------------- +# Replayable bodies keep retrying +# --------------------------------------------------------------------------- + + +def test_bytes_body_still_retries() -> None: + mock_http = MagicMock(spec=httpx.Client) + mock_http.request.side_effect = [ + httpx.ConnectError("boom"), + httpx.Response(200, request=httpx.Request("PUT", f"{BASE}/apis/test/v2/upload")), + ] + + client = NemoClient( + base_url=BASE, + http_client=mock_http, + retry=RetryPolicy(max_retries=3, backoff_base=0.0), + ) + resp = client.send(_upload_request(b"x" * PAYLOAD)) + + assert resp.http_response.status_code == 200 + assert mock_http.request.call_count == 2 + + +def test_bodyless_request_still_retries() -> None: + mock_http = MagicMock(spec=httpx.Client) + mock_http.request.side_effect = [ + httpx.Response(503, request=httpx.Request("GET", f"{BASE}/apis/test/v2/upload")), + httpx.Response(200, request=httpx.Request("GET", f"{BASE}/apis/test/v2/upload")), + ] + + client = NemoClient( + base_url=BASE, + http_client=mock_http, + retry=RetryPolicy(max_retries=3, backoff_base=0.0), + ) + request = PreparedRequest( + path_template="/apis/test/v2/upload", + path_params={}, + method="GET", + content=None, + content_type=None, + response_type=None, + ) + resp = client.send(request) + + assert resp.http_response.status_code == 200 + assert mock_http.request.call_count == 2 + + +def test_list_body_still_retries() -> None: + """httpx re-reads an in-memory sequence from the start, so it is replayable.""" + mock_http = MagicMock(spec=httpx.Client) + mock_http.request.side_effect = [ + httpx.ReadTimeout("timed out"), + httpx.Response(200, request=httpx.Request("PUT", f"{BASE}/apis/test/v2/upload")), + ] + + client = NemoClient( + base_url=BASE, + http_client=mock_http, + retry=RetryPolicy(max_retries=3, backoff_base=0.0), + ) + resp = client.send(_upload_request([b"x" * 1024, b"y" * 1024])) + + assert resp.http_response.status_code == 200 + assert mock_http.request.call_count == 2 + + +# --------------------------------------------------------------------------- +# A body that was never read is still safe to send again +# --------------------------------------------------------------------------- + + +def test_streaming_body_retries_when_the_connection_never_opened() -> None: + """Connect failures happen before httpx touches the body, so replay is safe.""" + mock_http = MagicMock(spec=httpx.Client) + mock_http.request.side_effect = [ + httpx.ConnectError("no route to host"), + httpx.Response(200, request=httpx.Request("PUT", f"{BASE}/apis/test/v2/upload")), + ] + + client = NemoClient( + base_url=BASE, + http_client=mock_http, + retry=RetryPolicy(max_retries=3, backoff_base=0.0), + ) + resp = client.with_headers({"Content-Length": str(PAYLOAD)}).send(_upload_request(_sync_chunks())) + + assert resp.http_response.status_code == 200 + assert mock_http.request.call_count == 2 + + +async def test_async_streaming_body_retries_when_the_connection_never_opened() -> None: + mock_http = MagicMock(spec=httpx.AsyncClient) + mock_http.request = AsyncMock( + side_effect=[ + httpx.ConnectTimeout("connect timed out"), + httpx.Response(200, request=httpx.Request("PUT", f"{BASE}/apis/test/v2/upload")), + ] + ) + + client = AsyncNemoClient( + base_url=BASE, + http_client=mock_http, + retry=RetryPolicy(max_retries=3, backoff_base=0.0), + ) + resp = await client.with_headers({"Content-Length": str(PAYLOAD)}).send(_upload_request(_chunks())) + + assert resp.http_response.status_code == 200 + assert mock_http.request.await_count == 2 + + +def test_streaming_body_is_not_retried_on_a_retryable_status() -> None: + """A response arrives only after the body is spent — there is nothing to resend.""" + mock_http = MagicMock(spec=httpx.Client) + mock_http.request.side_effect = [ + httpx.Response(503, request=httpx.Request("PUT", f"{BASE}/apis/test/v2/upload"), json={"detail": "down"}), + httpx.Response(200, request=httpx.Request("PUT", f"{BASE}/apis/test/v2/upload")), + ] + + client = NemoClient( + base_url=BASE, + http_client=mock_http, + retry=RetryPolicy(max_retries=3, backoff_base=0.0), + ) + + with pytest.raises(InternalServerError): + client.with_headers({"Content-Length": str(PAYLOAD)}).send(_upload_request(_sync_chunks())) + + assert mock_http.request.call_count == 1 + + +# --------------------------------------------------------------------------- +# The same, under the retry policy the SDK adapter actually installs +# --------------------------------------------------------------------------- + + +def _retry_me(status: int = 503) -> httpx.Response: + return httpx.Response( + status, + headers={"x-should-retry": "true"}, + request=httpx.Request("PUT", f"{BASE}/apis/test/v2/upload"), + json={"detail": "down"}, + ) + + +def test_one_shot_body_is_not_replayed_even_when_the_server_asks_for_a_retry() -> None: + """``x-should-retry: true`` cannot conjure back a body that is already spent. + + The spent-body check has to come before the header branches, not after, or a + server that sets this header would still get an exhausted iterator replayed + under the original Content-Length. + """ + mock_http = MagicMock(spec=httpx.Client) + mock_http.request.side_effect = [ + _retry_me(), + httpx.Response(200, request=httpx.Request("PUT", f"{BASE}/apis/test/v2/upload")), + ] + + client = NemoClient(base_url=BASE, http_client=mock_http, retry=ADAPTER_POLICY) + + with pytest.raises(InternalServerError): + client.with_headers({"Content-Length": str(PAYLOAD)}).send(_upload_request(_sync_chunks())) + + assert mock_http.request.call_count == 1 + + +def test_replayable_body_still_honors_the_server_retry_header() -> None: + """The spent-body check must not have swallowed header-driven retries wholesale.""" + mock_http = MagicMock(spec=httpx.Client) + mock_http.request.side_effect = [ + _retry_me(), + httpx.Response(200, request=httpx.Request("PUT", f"{BASE}/apis/test/v2/upload")), + ] + + client = NemoClient(base_url=BASE, http_client=mock_http, retry=ADAPTER_POLICY) + resp = client.send(_upload_request(b"x" * PAYLOAD)) + + assert resp.http_response.status_code == 200 + assert mock_http.request.call_count == 2 + + +# --------------------------------------------------------------------------- +# Declining a retry is worth a log line; everything else is not +# --------------------------------------------------------------------------- + + +def test_declined_retry_is_logged(caplog: pytest.LogCaptureFixture) -> None: + mock_http = MagicMock(spec=httpx.Client) + mock_http.request.side_effect = httpx.ReadTimeout("timed out") + + client = NemoClient( + base_url=BASE, + http_client=mock_http, + retry=RetryPolicy(max_retries=3, backoff_base=0.0), + ) + + with caplog.at_level(logging.INFO, logger=CLIENT_LOGGER), pytest.raises(NemoTransportError): + client.with_headers({"Content-Length": str(PAYLOAD)}).send(_upload_request(_sync_chunks())) + + assert "one-shot stream" in caplog.text + assert "PUT /apis/test/v2/upload" in caplog.text + + +def test_successful_upload_does_not_log_a_retry_decline(caplog: pytest.LogCaptureFixture) -> None: + """Nothing failed, so there was no retry to decline.""" + mock_http = MagicMock(spec=httpx.Client) + mock_http.request.return_value = httpx.Response(200, request=httpx.Request("PUT", f"{BASE}/apis/test/v2/upload")) + + client = NemoClient( + base_url=BASE, + http_client=mock_http, + retry=RetryPolicy(max_retries=3, backoff_base=0.0), + ) + + with caplog.at_level(logging.INFO, logger=CLIENT_LOGGER): + resp = client.with_headers({"Content-Length": str(PAYLOAD)}).send(_upload_request(_sync_chunks())) + + assert resp.http_response.status_code == 200 + assert caplog.records == [] + + +def test_non_retryable_status_does_not_log_a_retry_decline(caplog: pytest.LogCaptureFixture) -> None: + """A 404 stops the retry loop on its own merits — the body never came into it.""" + mock_http = MagicMock(spec=httpx.Client) + mock_http.request.return_value = httpx.Response( + 404, request=httpx.Request("PUT", f"{BASE}/apis/test/v2/upload"), json={"detail": "no such fileset"} + ) + + client = NemoClient( + base_url=BASE, + http_client=mock_http, + retry=RetryPolicy(max_retries=3, backoff_base=0.0), + ) + + with caplog.at_level(logging.INFO, logger=CLIENT_LOGGER), pytest.raises(NotFoundError): + client.with_headers({"Content-Length": str(PAYLOAD)}).send(_upload_request(_sync_chunks())) + + assert caplog.records == [] + + +# --------------------------------------------------------------------------- +# End to end over a real socket: the h11 failure mode itself +# --------------------------------------------------------------------------- + + +class _StallingServer: + """Accepts a request, drains the body, then stalls the first attempt.""" + + def __init__(self) -> None: + self.attempts: list[int] = [] + self._server: asyncio.AbstractServer | None = None + self._handlers: list[asyncio.Task] = [] + self.port = 0 + + async def __aenter__(self) -> _StallingServer: + self._server = await asyncio.start_server(self._handle, "127.0.0.1", 0) + self.port = self._server.sockets[0].getsockname()[1] + return self + + async def __aexit__(self, *exc: object) -> None: + assert self._server is not None + # ``wait_closed`` blocks on in-flight handlers, and the stalling one sleeps + # far longer than the test needs. Cancel them so teardown is immediate. + for handler in self._handlers: + handler.cancel() + self._server.close() + await self._server.wait_closed() + + async def _handle(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None: + task = asyncio.current_task() + if task is not None: + self._handlers.append(task) + attempt = len(self.attempts) + 1 + try: + head = await reader.readuntil(b"\r\n\r\n") + declared = 0 + for line in head.split(b"\r\n"): + if line.lower().startswith(b"content-length:"): + declared = int(line.split(b":", 1)[1]) + remaining = declared + while remaining > 0: + data = await reader.read(min(65536, remaining)) + if not data: + break + remaining -= len(data) + self.attempts.append(declared - remaining) + if attempt == 1: + # Stall past the client read timeout, as a Files service writing + # multiple GB to storage would before it answers. + await asyncio.sleep(STALL_SECONDS) + body = b"{}" + writer.write(b"HTTP/1.1 200 OK\r\nContent-Length: %d\r\n\r\n" % len(body) + body) + await writer.drain() + except (asyncio.CancelledError, ConnectionError, asyncio.IncompleteReadError): + pass + finally: + writer.close() + + +async def test_read_timeout_does_not_become_a_content_length_error() -> None: + """The upload's real failure must survive, not be masked by an h11 abort.""" + async with _StallingServer() as server: + base = f"http://127.0.0.1:{server.port}" + async with httpx.AsyncClient(base_url=base, timeout=httpx.Timeout(5.0, read=0.5)) as http: + client = AsyncNemoClient( + base_url=base, + http_client=http, + retry=RetryPolicy(max_retries=2, backoff_base=0.0), + ) + + with pytest.raises(NemoTransportError) as excinfo: + await client.with_headers({"Content-Length": str(PAYLOAD)}).send(_upload_request(_chunks())) + + assert isinstance(excinfo.value.error, httpx.ReadTimeout) + # One attempt only. A second would have re-sent an exhausted iterator under + # the original Content-Length and raised h11's LocalProtocolError instead. + assert server.attempts == [PAYLOAD] diff --git a/packages/nmp_customization_common/src/nmp/customization_common/tasks/file_io/run.py b/packages/nmp_customization_common/src/nmp/customization_common/tasks/file_io/run.py index 02b5c7051b..7c3951daf7 100644 --- a/packages/nmp_customization_common/src/nmp/customization_common/tasks/file_io/run.py +++ b/packages/nmp_customization_common/src/nmp/customization_common/tasks/file_io/run.py @@ -24,23 +24,22 @@ import httpx from nemo_platform import ( - APIConnectionError, - APITimeoutError, - InternalServerError, NeMoPlatform, NotFoundError, ) -from nemo_platform.types.files.fileset_file import FilesetFile from nemo_platform_plugin.client.adapter import client_from_platform from nemo_platform_plugin.client.errors import ( ConflictError, + NemoTransportError, + RateLimitError, ) from nemo_platform_plugin.client.errors import ( InternalServerError as ClientInternalServerError, ) from nemo_platform_plugin.client.types import RetryPolicy from nemo_platform_plugin.files.client import FilesClient -from nemo_platform_plugin.files.types import CreateFilesetRequest, UpdateFilesetRequest +from nemo_platform_plugin.files.metadata import FilesetMetadata +from nemo_platform_plugin.files.types import CreateFilesetRequest, FilesetFileOutput, UpdateFilesetRequest from nmp.common.jobs.schemas import PlatformJobStatus from nmp.common.sdk_factory import get_task_sdk from nmp.customization_common.schemas.file_io import ( @@ -84,12 +83,23 @@ INITIAL_BACKOFF_SECONDS = 1.0 MAX_BACKOFF_SECONDS = 30.0 +# Transient failures worth another attempt, for every Files operation this task +# runs. Uploads stream a one-shot body, so once it is on the wire the typed client +# can no longer retry them itself — this is the layer that rebuilds the request +# from the source file, so it has to catch what the client raises. Those are the +# Nemo* entries: they wrap transport failures, 5xx responses (raise_for_status +# maps every 5xx to InternalServerError) and 429 respectively, and are not httpx +# types. Anything reaching this task through the typed client arrives as one of +# them, never as the bare httpx error underneath. TRANSIENT_FILESYSTEM_EXCEPTIONS = ( httpx.TimeoutException, httpx.ConnectError, httpx.ReadTimeout, httpx.RemoteProtocolError, httpx.ReadError, + NemoTransportError, + ClientInternalServerError, + RateLimitError, ) @@ -109,8 +119,8 @@ def __init__( self.job_ctx = job_ctx self.service_source = service_source - def list_fileset_files(self, fileset: FileSetRef) -> list[FilesetFile]: - """List files in a FileSet. Returns a list of ``FilesetFile`` objects.""" + def list_fileset_files(self, fileset: FileSetRef) -> list[FilesetFileOutput]: + """List files in a FileSet. Returns a list of ``FilesetFileOutput`` objects.""" try: with sdk_error_handler(FileDownloadError, f"list files in fileset {fileset}", passthrough=(NotFoundError,)): response = self.sdk.with_options(timeout=LIST_FILES_TIMEOUT).files.list( @@ -272,17 +282,9 @@ def create_fileset(self, fileset: FileSetRef, metadata: dict | None = None) -> N @retry( stop=stop_after_attempt(MAX_RETRIES), wait=wait_exponential(multiplier=2, min=INITIAL_BACKOFF_SECONDS, max=MAX_BACKOFF_SECONDS), - retry=retry_if_exception_type( - ( - InternalServerError, - APITimeoutError, - APIConnectionError, - ClientInternalServerError, - httpx.TimeoutException, - httpx.ConnectError, - ) - ), + retry=retry_if_exception_type(TRANSIENT_FILESYSTEM_EXCEPTIONS), reraise=True, + before_sleep=before_sleep_log(logger, logging.WARNING), ) def _create_fileset_with_retry(self, fileset: FileSetRef, metadata: dict | None = None) -> None: """Internal method with retry logic for creating a FileSet.""" @@ -295,7 +297,7 @@ def _create_fileset_with_retry(self, fileset: FileSetRef, metadata: dict | None "custom_fields": {"service_source": self.service_source}, } if metadata is not None: - body_kwargs["metadata"] = metadata + body_kwargs["metadata"] = FilesetMetadata.model_validate(metadata) result = files.create_fileset(workspace=fileset.workspace, body=CreateFilesetRequest(**body_kwargs)).data() logger.info(f"Created FileSet: {result.workspace}/{result.name}") except ConflictError: @@ -305,7 +307,7 @@ def _create_fileset_with_retry(self, fileset: FileSetRef, metadata: dict | None files.update_fileset( workspace=workspace, name=fileset.name, - body=UpdateFilesetRequest(metadata=metadata), + body=UpdateFilesetRequest(metadata=FilesetMetadata.model_validate(metadata)), ) logger.info(f"Patched existing FileSet metadata: {workspace}/{fileset.name}") except Exception as e: diff --git a/packages/nmp_customization_common/tests/tasks/test_file_io.py b/packages/nmp_customization_common/tests/tasks/test_file_io.py index 81960e94e7..1327bc9a2f 100644 --- a/packages/nmp_customization_common/tests/tasks/test_file_io.py +++ b/packages/nmp_customization_common/tests/tasks/test_file_io.py @@ -9,6 +9,7 @@ from pathlib import Path from unittest.mock import MagicMock, patch +import httpx import pytest @@ -197,3 +198,114 @@ def test_empty_fileset_returns_zero_stats_without_downloading(self, tmp_path: Pa assert stats.files_downloaded == 0 assert stats.total_bytes == 0 sdk.files.download.assert_not_called() + + +@pytest.fixture +def no_retry_backoff(monkeypatch: pytest.MonkeyPatch) -> None: + """Skip tenacity's exponential backoff so retry tests don't sleep for seconds. + + ``tenacity.nap.sleep`` resolves ``time.sleep`` per call, so patching it here + takes effect even though the retry policy was bound at decoration time. + """ + monkeypatch.setattr("tenacity.nap.time.sleep", lambda _seconds: None) + + +class TestCreateFilesetRetry: + """Create goes through the typed client, so it must catch the typed client's errors.""" + + @patch("nmp.customization_common.tasks.file_io.run.client_from_platform") + def test_retries_transport_error_wrapped_by_the_client(self, mock_cfp, no_retry_backoff) -> None: + from nemo_platform_plugin.client.errors import NemoTransportError + from nmp.customization_common.schemas.file_io import FileSetRef + + mock_fc = MagicMock() + mock_fc.with_options.return_value = mock_fc + mock_fc.create_fileset.side_effect = [NemoTransportError(httpx.ConnectError("refused")), MagicMock()] + mock_cfp.return_value = mock_fc + runner = _make_runner(_make_sdk()) + + runner.create_fileset(FileSetRef(workspace="default", name="models")) + + assert mock_fc.create_fileset.call_count == 2 + + @patch("nmp.customization_common.tasks.file_io.run.client_from_platform") + def test_retries_rate_limit_wrapped_by_the_client(self, mock_cfp, no_retry_backoff) -> None: + from nemo_platform_plugin.client.errors import RateLimitError + from nmp.customization_common.schemas.file_io import FileSetRef + + response = httpx.Response(429, request=httpx.Request("POST", "http://test/filesets"), json={"detail": "slow"}) + mock_fc = MagicMock() + mock_fc.with_options.return_value = mock_fc + mock_fc.create_fileset.side_effect = [RateLimitError(response), MagicMock()] + mock_cfp.return_value = mock_fc + runner = _make_runner(_make_sdk()) + + runner.create_fileset(FileSetRef(workspace="default", name="models")) + + assert mock_fc.create_fileset.call_count == 2 + + +class TestUploadRetry: + """The task layer owns upload retries. + + The typed client cannot retry a streaming upload — the body is a one-shot + iterator, and replaying it under the original Content-Length makes h11 abort + the request. So the client raises through, and this layer, which rebuilds the + request from the source file on every attempt, is where the retry belongs. + """ + + def test_retries_transport_error_wrapped_by_the_client(self, tmp_path: Path, no_retry_backoff) -> None: + from nemo_platform_plugin.client.errors import NemoTransportError + from nmp.customization_common.schemas.file_io import FileSetRef + + src = _make_dir(tmp_path) + sdk = _make_sdk() + sdk.files.upload.side_effect = [NemoTransportError(httpx.ReadTimeout("timed out")), None] + runner = _make_runner(sdk) + + runner.upload_fileset(FileSetRef(workspace="default", name="models"), src.resolve()) + + assert sdk.files.upload.call_count == 2 + + def test_retries_server_error_wrapped_by_the_client(self, tmp_path: Path, no_retry_backoff) -> None: + from nemo_platform_plugin.client.errors import InternalServerError + from nmp.customization_common.schemas.file_io import FileSetRef + + src = _make_dir(tmp_path) + sdk = _make_sdk() + response = httpx.Response(503, request=httpx.Request("PUT", "http://test/upload"), json={"detail": "down"}) + sdk.files.upload.side_effect = [InternalServerError(response), None] + runner = _make_runner(sdk) + + runner.upload_fileset(FileSetRef(workspace="default", name="models"), src.resolve()) + + assert sdk.files.upload.call_count == 2 + + def test_retries_rate_limit_wrapped_by_the_client(self, tmp_path: Path, no_retry_backoff) -> None: + """429 is in the client's retryable statuses, but it cannot act on it here.""" + from nemo_platform_plugin.client.errors import RateLimitError + from nmp.customization_common.schemas.file_io import FileSetRef + + src = _make_dir(tmp_path) + sdk = _make_sdk() + response = httpx.Response(429, request=httpx.Request("PUT", "http://test/upload"), json={"detail": "slow down"}) + sdk.files.upload.side_effect = [RateLimitError(response), None] + runner = _make_runner(sdk) + + runner.upload_fileset(FileSetRef(workspace="default", name="models"), src.resolve()) + + assert sdk.files.upload.call_count == 2 + + def test_gives_up_as_a_file_upload_error(self, tmp_path: Path, no_retry_backoff) -> None: + from nemo_platform_plugin.client.errors import NemoTransportError + from nmp.customization_common.schemas.file_io import FileSetRef, FileUploadError + + src = _make_dir(tmp_path) + sdk = _make_sdk() + sdk.files.upload.side_effect = NemoTransportError(httpx.ReadTimeout("timed out")) + runner = _make_runner(sdk) + + with pytest.raises(FileUploadError): + runner.upload_fileset(FileSetRef(workspace="default", name="models"), src.resolve()) + + assert sdk.files.upload.call_count == 3 diff --git a/sdk/python/nemo-platform/src/nemo_platform/filesets/filesystem/filesystem.py b/sdk/python/nemo-platform/src/nemo_platform/filesets/filesystem/filesystem.py index 9862c0def3..ac400f7dc9 100644 --- a/sdk/python/nemo-platform/src/nemo_platform/filesets/filesystem/filesystem.py +++ b/sdk/python/nemo-platform/src/nemo_platform/filesets/filesystem/filesystem.py @@ -358,24 +358,24 @@ def _ensure_async(client: FilesClient | AsyncFilesClient) -> AsyncFilesClient: import httpx + # A timeout lives in two layers, so mirror each from its own source: the + # transport carries the client-level default, and ``_timeout`` the + # per-request override that ``send`` puts on every request. Leave the + # transport's unset and httpx falls back to its own 5s, which a multi-GB + # upload blows through waiting for the server to commit the body to storage. asgi_app = getattr(client._http, "asgi_app", None) - http_client = ( - httpx.AsyncClient( - transport=httpx.ASGITransport(app=asgi_app), - base_url=client.base_url, - headers=dict(client._default_headers) if client._default_headers else None, - ) - if asgi_app is not None - else httpx.AsyncClient( - base_url=client.base_url, - headers=dict(client._default_headers) if client._default_headers else None, - ) + http_client = httpx.AsyncClient( + transport=httpx.ASGITransport(app=asgi_app) if asgi_app is not None else None, + base_url=client.base_url, + headers=dict(client._default_headers) if client._default_headers else None, + timeout=client._http.timeout, ) return AsyncFilesClient( base_url=client.base_url, workspace=client.workspace, auth=client._auth, default_headers=client._default_headers or None, + timeout=client._timeout, retry=client._retry, http_client=http_client, url_resolver=client._url_resolver,