[tune][release] Upgrade tune_torch_benchmark to v2 #56804
Merged
Conversation
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
Signed-off-by: Lehui Liu <lehui@anyscale.com>
Signed-off-by: Lehui Liu <lehui@anyscale.com>
![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
This comment was marked as outdated.
This comment was marked as outdated.
Sorry, something went wrong.
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Code Review
This pull request successfully upgrades the tune_torch_benchmark to use the Ray Train V2 API. The changes correctly adapt the benchmark to the new V2 patterns, such as using a train_driver_fn for tuning and leveraging TuneReportCallback. The refactoring to move train_loop to the module level improves code clarity. Additionally, the test configuration in release_tests.yaml has been updated to run faster and enable the V2 API, which is a sensible adjustment for release testing.
I have one minor suggestion to improve the robustness of the tune_torch function against a potential TypeError if it's called with a None config, as allowed by its signature.
release/air_tests/air_benchmarks/workloads/tune_torch_benchmark.py
Outdated
Show resolved
Hide resolved
Signed-off-by: Lehui Liu <lehui@anyscale.com>
This comment was marked as outdated.
This comment was marked as outdated.
Sorry, something went wrong.
Signed-off-by: Lehui Liu <lehui@anyscale.com>
justinvyu
approved these changes
Sep 23, 2025
release/air_tests/air_benchmarks/workloads/tune_torch_benchmark.py
Outdated
Show resolved
Hide resolved
Signed-off-by: Lehui Liu <lehui@anyscale.com>
elliot-barn
pushed a commit
that referenced
this pull request
Sep 27, 2025
1. Update to Train V2 Train+Tune integration API. 2. Pass in the TuneReportCallback for the trainer that used in Tune for reported results. 3. Reduced the number of runs / trials to make the test run faster --------- Signed-off-by: Lehui Liu <lehui@anyscale.com> Signed-off-by: elliot-barn <elliot.barnwell@anyscale.com>
dstrodtman
pushed a commit
that referenced
this pull request
Oct 6, 2025
1. Update to Train V2 Train+Tune integration API. 2. Pass in the TuneReportCallback for the trainer that used in Tune for reported results. 3. Reduced the number of runs / trials to make the test run faster --------- Signed-off-by: Lehui Liu <lehui@anyscale.com> Signed-off-by: Douglas Strodtman <douglas@anyscale.com>
justinyeh1995
pushed a commit
to justinyeh1995/ray
that referenced
this pull request
Oct 20, 2025
1. Update to Train V2 Train+Tune integration API. 2. Pass in the TuneReportCallback for the trainer that used in Tune for reported results. 3. Reduced the number of runs / trials to make the test run faster --------- Signed-off-by: Lehui Liu <lehui@anyscale.com>
landscapepainter
pushed a commit
to landscapepainter/ray
that referenced
this pull request
Nov 17, 2025
1. Update to Train V2 Train+Tune integration API. 2. Pass in the TuneReportCallback for the trainer that used in Tune for reported results. 3. Reduced the number of runs / trials to make the test run faster --------- Signed-off-by: Lehui Liu <lehui@anyscale.com>
Future-Outlier
pushed a commit
to Future-Outlier/ray
that referenced
this pull request
Dec 7, 2025
1. Update to Train V2 Train+Tune integration API. 2. Pass in the TuneReportCallback for the trainer that used in Tune for reported results. 3. Reduced the number of runs / trials to make the test run faster --------- Signed-off-by: Lehui Liu <lehui@anyscale.com> Signed-off-by: Future-Outlier <eric901201@gmail.com>
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 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.
Why are these changes needed?
Related issue number
Checks
git commit -s) in this PR.scripts/format.shto lint the changes in this PR.method in Tune, I've added it in
doc/source/tune/api/under thecorresponding
.rstfile.