[data][train] Refactor call_with_retry into shared library and use it to retry checkpoint upload by TimothySeah · Pull Request #56608 · ray-project/ray

TimothySeah · 2025-09-17T02:40:18Z

Summary

This PR moves call_with_retry from ray/data/_internal to ray/_private so that it can be used in other libraries like Ray Train.

It also adds a new retry decorator that wraps around call_with_retry. Note that I had to remove * from call_with_retry's arguments to get the decorator to work on Python object methods because Python passes self as one of the *args.

Finally, it uses this decorator to retry checkpoint uploads.

Testing

Unit tests

Signed-off-by: Timothy Seah <tseah@anyscale.com>

gemini-code-assist

Code Review

This pull request effectively refactors the retry logic into a shared utility call_with_retry and introduces a convenient @retry decorator. The changes improve code modularity and reusability. The new decorator is correctly applied to add resilience to the checkpoint uploading process in Ray Train. I have a couple of suggestions to improve the logging clarity and documentation of the new retry utility.

python/ray/_private/retry.py

edoakes · 2025-09-17T13:20:07Z

utilities used by multiple libraries should live in ray._common

please move it there and add standalone tests for the utility

justinvyu

Nice!

justinvyu · 2025-09-17T16:23:59Z

python/ray/train/v2/_internal/execution/context.py

            )
        )

+    @retry(description="upload checkpoint", max_attempts=3)


should we make this 3 configurable? also, I think we should match specific "upload errors" similar to data's whitelist:

https://github.com/anyscale/rayturbo/blob/788a223b6c933a303de0212b34883fdf0a1f4977/python/ray/data/context.py#L162

Otherwise we'll retry on the other "unretryable" errors that we raise explicitly already in _upload_checkpoint.

Added TODO for retry configurability: could be a good extension to #55861.

Added COMMON_RETRYABLE_TOKENS lmk if that's fine.

python/ray/train/v2/_internal/execution/context.py

Signed-off-by: Timothy Seah <tseah@anyscale.com>

…point-upload

python/ray/_common/retry.py

Signed-off-by: Timothy Seah <tseah@anyscale.com>

python/ray/_common/retry.py

Signed-off-by: Timothy Seah <tseah@anyscale.com>

…point-upload

Signed-off-by: Timothy Seah <tseah@anyscale.com>

… to retry checkpoint upload (ray-project#56608) This PR moves `call_with_retry` from `ray/data/_internal` to `ray/_private` so that it can be used in other libraries like Ray Train. It also adds a new `retry` decorator that wraps around `call_with_retry`. Note that I had to remove `*` from `call_with_retry`'s arguments to get the decorator to work on Python object methods because Python passes `self` as one of the `*args`. --------- Signed-off-by: Timothy Seah <tseah@anyscale.com> Signed-off-by: zac <zac@anyscale.com>

… to retry checkpoint upload (#56608) This PR moves `call_with_retry` from `ray/data/_internal` to `ray/_private` so that it can be used in other libraries like Ray Train. It also adds a new `retry` decorator that wraps around `call_with_retry`. Note that I had to remove `*` from `call_with_retry`'s arguments to get the decorator to work on Python object methods because Python passes `self` as one of the `*args`. --------- Signed-off-by: Timothy Seah <tseah@anyscale.com> Signed-off-by: elliot-barn <elliot.barnwell@anyscale.com>

… to retry checkpoint upload (ray-project#56608) This PR moves `call_with_retry` from `ray/data/_internal` to `ray/_private` so that it can be used in other libraries like Ray Train. It also adds a new `retry` decorator that wraps around `call_with_retry`. Note that I had to remove `*` from `call_with_retry`'s arguments to get the decorator to work on Python object methods because Python passes `self` as one of the `*args`. --------- Signed-off-by: Timothy Seah <tseah@anyscale.com> Signed-off-by: Marco Stephan <marco@magic.dev>

… to retry checkpoint upload (#56608) This PR moves `call_with_retry` from `ray/data/_internal` to `ray/_private` so that it can be used in other libraries like Ray Train. It also adds a new `retry` decorator that wraps around `call_with_retry`. Note that I had to remove `*` from `call_with_retry`'s arguments to get the decorator to work on Python object methods because Python passes `self` as one of the `*args`. --------- Signed-off-by: Timothy Seah <tseah@anyscale.com> Signed-off-by: elliot-barn <elliot.barnwell@anyscale.com>

… to retry checkpoint upload (#56608) This PR moves `call_with_retry` from `ray/data/_internal` to `ray/_private` so that it can be used in other libraries like Ray Train. It also adds a new `retry` decorator that wraps around `call_with_retry`. Note that I had to remove `*` from `call_with_retry`'s arguments to get the decorator to work on Python object methods because Python passes `self` as one of the `*args`. --------- Signed-off-by: Timothy Seah <tseah@anyscale.com> Signed-off-by: Douglas Strodtman <douglas@anyscale.com>

… to retry checkpoint upload (ray-project#56608) This PR moves `call_with_retry` from `ray/data/_internal` to `ray/_private` so that it can be used in other libraries like Ray Train. It also adds a new `retry` decorator that wraps around `call_with_retry`. Note that I had to remove `*` from `call_with_retry`'s arguments to get the decorator to work on Python object methods because Python passes `self` as one of the `*args`. --------- Signed-off-by: Timothy Seah <tseah@anyscale.com>

… to retry checkpoint upload (ray-project#56608) This PR moves `call_with_retry` from `ray/data/_internal` to `ray/_private` so that it can be used in other libraries like Ray Train. It also adds a new `retry` decorator that wraps around `call_with_retry`. Note that I had to remove `*` from `call_with_retry`'s arguments to get the decorator to work on Python object methods because Python passes `self` as one of the `*args`. --------- Signed-off-by: Timothy Seah <tseah@anyscale.com> Signed-off-by: Future-Outlier <eric901201@gmail.com>

[data][train] Refactor call_with_retry into shared library

c80ebff

Signed-off-by: Timothy Seah <tseah@anyscale.com>

TimothySeah requested review from a team as code owners September 17, 2025 02:40

gemini-code-assist bot reviewed Sep 17, 2025

View reviewed changes

python/ray/_private/retry.py Show resolved Hide resolved

python/ray/_private/retry.py Outdated Show resolved Hide resolved

ray-gardener bot added train Ray Train Related Issue core Issues that should be addressed in Ray Core data Ray Data-related issues labels Sep 17, 2025

justinvyu reviewed Sep 17, 2025

View reviewed changes

address pr comments

6d1de38

Signed-off-by: Timothy Seah <tseah@anyscale.com>

TimothySeah requested review from edoakes and jjyao as code owners September 18, 2025 00:49

add todo

eb6630f

Signed-off-by: Timothy Seah <tseah@anyscale.com>

TimothySeah requested a review from justinvyu September 18, 2025 00:53

Merge remote-tracking branch 'upstream/master' into tseah/retry-check…

2e6ad07

…point-upload

edoakes added the go add ONLY when ready to merge, run all tests label Sep 19, 2025

edoakes approved these changes Sep 19, 2025

View reviewed changes

justinvyu approved these changes Sep 20, 2025

View reviewed changes

python/ray/_common/retry.py Outdated Show resolved Hide resolved

moved retryable tokens to train with aws prefixes

f981df9

Signed-off-by: Timothy Seah <tseah@anyscale.com>

This comment was marked as outdated.

Sign in to view

fix assertion

0323ee6

Signed-off-by: Timothy Seah <tseah@anyscale.com>

alexeykudinkin reviewed Sep 22, 2025

View reviewed changes

python/ray/_common/retry.py Outdated Show resolved Hide resolved

python/ray/_common/retry.py Outdated Show resolved Hide resolved

convert exception to str once + use 20% jitter

f701580

Signed-off-by: Timothy Seah <tseah@anyscale.com>

TimothySeah requested a review from alexeykudinkin September 22, 2025 21:05

alexeykudinkin approved these changes Sep 22, 2025

View reviewed changes

TimothySeah added 2 commits September 22, 2025 18:28

Merge remote-tracking branch 'upstream/master' into tseah/retry-check…

17154f5

…point-upload

attempt to fix timeout by making average retry wait the same as before

d13a897

Signed-off-by: Timothy Seah <tseah@anyscale.com>

This comment was marked as outdated.

Sign in to view

fix unit test

ce89a27

Signed-off-by: Timothy Seah <tseah@anyscale.com>

justinvyu merged commit c4be355 into ray-project:master Sep 23, 2025
6 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[data][train] Refactor call_with_retry into shared library and use it to retry checkpoint upload#56608

[data][train] Refactor call_with_retry into shared library and use it to retry checkpoint upload#56608
justinvyu merged 10 commits intoray-project:masterfrom
TimothySeah:tseah/retry-checkpoint-upload

TimothySeah commented Sep 17, 2025 •

edited

Loading

Uh oh!

gemini-code-assist bot left a comment

Uh oh!

Uh oh!

Uh oh!

edoakes commented Sep 17, 2025

Uh oh!

justinvyu left a comment

Uh oh!

justinvyu Sep 17, 2025

Uh oh!

TimothySeah Sep 18, 2025

Uh oh!

Uh oh!

Uh oh!

This comment was marked as outdated.

Uh oh!

Uh oh!

Uh oh!

This comment was marked as outdated.

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Conversation

TimothySeah commented Sep 17, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Testing

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

edoakes commented Sep 17, 2025

Uh oh!

justinvyu left a comment

Choose a reason for hiding this comment

Uh oh!

justinvyu Sep 17, 2025

Choose a reason for hiding this comment

Uh oh!

TimothySeah Sep 18, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

This comment was marked as outdated.

Uh oh!

Uh oh!

Uh oh!

This comment was marked as outdated.

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

TimothySeah commented Sep 17, 2025 •

edited

Loading