fix(client): FileSet upload retries and timeout propagation - #1125
Merged
Conversation
albcui
force-pushed
the
albcui/NVBUG6562815-fix-fileset-upload-retry-body
branch
from
August 6, 2026 14:24
df72839 to
b93e86a
Compare
albcui
marked this pull request as ready for review
August 6, 2026 14:25
Contributor
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughThe change propagates timeouts through platform and filesystem clients, restricts retries for consumed request bodies, and updates file operations to use typed outputs, validated metadata, and shared transient exceptions. ChangesClient reliability and file operations
Sequence Diagram(s)sequenceDiagram
participant Caller
participant NemoClient
participant RequestBody
participant Server
Caller->>NemoClient: send request
NemoClient->>RequestBody: inspect replayability
NemoClient->>Server: transmit request
Server-->>NemoClient: response or transport failure
NemoClient->>NemoClient: retry only when body state permits
NemoClient-->>Caller: response or NemoTransportError
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1📝 Generate docstrings 💡
🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with 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.
Inline comments:
In `@packages/nemo_platform_plugin/src/nemo_platform_plugin/client/client.py`:
- Around line 245-251: Update the retry-suppression log condition near
_should_retry so it runs only when the response would otherwise be retryable
with replayable=True and produce a backoff, while the actual request is
non-replayable, backoff is None, and retries remain. Do not log for successful
or otherwise non-retryable responses; preserve the existing retry and return
behavior.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 42fd85f4-47aa-4dcb-ab0f-4f0f27ff9e0d
⛔ Files ignored due to path filters (1)
sdk/python/nemo-platform/src/nemo_platform/filesets/filesystem/filesystem.pyis excluded by!sdk/**
📒 Files selected for processing (9)
packages/filesets/src/filesets/filesystem/filesystem.pypackages/filesets/tests/test_filesystem_client.pypackages/nemo_platform_plugin/src/nemo_platform_plugin/client/adapter.pypackages/nemo_platform_plugin/src/nemo_platform_plugin/client/client.pypackages/nemo_platform_plugin/src/nemo_platform_plugin/client/types.pypackages/nemo_platform_plugin/tests/client/test_adapter.pypackages/nemo_platform_plugin/tests/client/test_retry_streaming_body.pypackages/nmp_customization_common/src/nmp/customization_common/tasks/file_io/run.pypackages/nmp_customization_common/tests/tasks/test_file_io.py
Contributor
|
albcui
changed the title
… through Retrying a streaming upload re-sent an exhausted iterator under the original Content-Length, so h11 aborted with "Too little data for declared Content-Length" and masked the real failure. Retry such a body only while it is still untouched -- that is, on a connection-establishment failure, which happens before httpx reads any of it -- and surface anything later to the caller. A retryable status code is never replayed either, since the body is spent by the time a response arrives. file_io's tenacity layer rebuilds the request from the source file on every attempt, so it is where an upload retry belongs. Teach it to catch what the client actually raises: NemoTransportError, InternalServerError and RateLimitError are client types, not httpx ones, so none of them were being retried. Also carry the platform's timeout into the typed client and onto the async transport FilesetFileSystem builds. Both silently fell back to their own defaults -- the latter to httpx's 5s, which a multi-GB upload blows through waiting for the server to commit the body to storage. Drive-by, so the ty pre-commit hook passes on the touched files: correct list_fileset_files' return annotation to the FilesetFileOutput it really returns, validate fileset metadata into FilesetMetadata explicitly, and add ty suppressions next to the existing mypy ones in test_adapter. NVBUG6562815 Signed-off-by: Albert Cui <albcui@nvidia.com>
… see _create_fileset_with_retry talks only to the typed client, which wraps everything it raises. Five of the six entries in its retry tuple could never fire: the two bare httpx types arrive as NemoTransportError, and the three Stainless SDK types belong to a path this method never takes. Only ClientInternalServerError was live. So a connection refusal or timeout against the Files service -- the case the retry exists for -- sailed past tenacity on the first attempt and came out of sdk_error_handler as a terminal FileUploadError. Same root cause as the upload retry, second call site. Share TRANSIENT_FILESYSTEM_EXCEPTIONS instead, which adds NemoTransportError and RateLimitError and drops the dead entries. Ruff then found the three SDK imports unused across the whole module, confirming they were unreachable rather than merely redundant here. Signed-off-by: Albert Cui <albcui@nvidia.com>
_should_retry_request logged "not retrying: one-shot stream" on every outcome it did not retry, which included every successful upload -- one misleading line per file in the customizer's job logs. Ask _should_retry again as if the body were replayable to tell that case apart from a success, a 404, or an exhausted attempt budget. Drop bytearray and memoryview from _is_replayable. httpx treats anything that is not bytes/str as an iterable of chunks, and iterating either of those yields int, so it rejects both on the first attempt regardless of what the predicate claims -- neither could ever have reached a retry. Cover the interaction the spent-body check has with the retry-decision headers client_from_platform turns on: a server sending x-should-retry: true cannot conjure back a body that is already spent, and a replayable body must still honour it. Neither had a test, and moving the check below that branch silently breaks the first. Also drop an isinstance guard in the adapter that can never be false, give _create_fileset_with_retry the before_sleep_log its two sibling retries already have, and validate create's fileset metadata into FilesetMetadata the way the update path does. NVBUG6562815 Signed-off-by: Albert Cui <albcui@nvidia.com>
_should_retry and _should_retry_request only ever ran as a pair: the first decided, the second added the body's say in it. The split was an artefact of layering one-shot-body handling onto the existing decision logic, and it cost a second _should_retry call just to work out whether the body was the reason a retry stopped. Let the policy decide first and give the body a veto at the end. Reaching the tail already means a retry was due, so the veto is the only thing that can stop it there -- which is exactly when the decline is worth logging. The replayable keyword and the repeat call both fall out. The x-should-retry branch no longer returns early either. It and the status-code path applied the same retry-after logic, so they now share a tail, which is what lets the veto cover both without being written twice. No behaviour change: the same 207 tests pass, and both mutations the streaming-body tests were written to catch still fail against the consolidated function. Signed-off-by: Albert Cui <albcui@nvidia.com>
albcui
force-pushed
the
albcui/NVBUG6562815-fix-fileset-upload-retry-body
branch
from
August 6, 2026 16:03
28ff4c2 to
ab3ea2b
Compare
Contributor
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
albcui
changed the title
albcui
requested review from
matthewgrossman and
soluwalana
and removed request for
soluwalana
August 6, 2026 20:32
ironcommit
approved these changes
Aug 6, 2026
ironcommit
left a comment
ironcommit
left a comment
Contributor
There was a problem hiding this comment.
LGTM. This works around our very messy transport system. But it will do.
matthewgrossman
approved these changes
Aug 6, 2026
Signed-off-by: Albert Cui <albcui@nvidia.com>
albcui
enabled auto-merge
August 6, 2026 22:24
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
It was observed that a FileSet upload that hit a transient failure failed with
h11._util.LocalProtocolError: Too little data for declared Content-Lengtherror. However, this is a bit of a red herring, because it hides the underlying issue, which is a bit gnarly, so let's dig a bit deeper.Root cause
The core crux of the issue is in
filesystem.py'sFilesetFileSystem._put_file, which implementsfsspec.The
stream_file()is a generator object that gets called and stored once inPreparedRequest.contentand reused by every pass ofNemoClient's retry loop. If any transient failures occur, then the sameContent-Lengthfollowed by zero bytes (the generator is already exhausted) is sent, which breaks the HTTP/1.1 protocol, and we get ahttpx.LocalProtocolError, which subclassesTransportError, which gets retried until attempts ran out. This is bad because we actually lose the information for the original error.So what could've caused the original failure?
Let's start with entrypoint from Customizer's
file_iotaks:The expected behaviour is that
UPLOAD_TIMEOUTwould propagate down tofiles.upload, which should extend the write timeout to a generous 600s. However, this is not the case. What actually happened was:with_options(...)created a newNeMoPlatformclient with the new timeout based onUPLOAD_TIMEOUT, but keeps the originalhttpx.Client, which has its own timeout (60s). So now there are two timeouts:NeMoPlatform.timeout = Timeout(30, write=600, read=300)NeMoPlatform._client.timeout = Timeout(60, ...)client_from_platformshares theNeMoPlatform._client. The sharing is not wrong by itself, it's the fact that we didn't propagate theNeMoPlatform.timeoutin the first place.FilesetFileSystem._ensure_asyncbuilt a brand newhttps.AsyncClientwith no timeout at all, which defaults to5s.Essentially,
FilesetFileSystem._put_file's httpx client had a timeout of 5s (on both writes and reads). This could error for both scenarios:The fix involved changes in multiple layers of the code:
client/client.py-- don't retry requests with generator bodies, at least not at this layer of the codetasks/file_io/run.py-- retry here instead, since it can rebuild the generatorclient/adapter.py+filesystem/filesystem.py-- propagate the timeout properlyclient.with_options(timeout=platform.timeout)ensure_asyncpropagates the timeout onto both the newhttps.AsyncClientas well as theAsyncFilesClientSummary by CodeRabbit
Bug Fixes
Bug Fixes
Documentation