[train] Enable v2 for ray/train/tests by justinvyu · Pull Request #56868 · ray-project/ray

justinvyu · 2025-09-24T01:01:45Z

Note

Enable Train v2 across CI with new CPU/GPU jobs, migrate tests to v2 with env flags and tag updates, remove legacy Tune examples/utilities, and fix docs/BUILD configurations.

CI (Buildkite):
- Split Train tests into v1 and v2, add dedicated GPU job for train_v2_gpu.
- Standardize tags (drop gpu_only, use gpu), adjust include/except tag filters.
Tests (BUILD and test code):
- Gate tests with RAY_TRAIN_V2_ENABLED (on/off per test); retag v2 tests (train_v2, train_v2_gpu).
- Replace deprecated imports/usages; simplify some assertions to just run fit().
- Remove legacy tests/examples (e.g., Tune Torch/TensorFlow tuning examples, datasets_train, utils), and adjust v2 HF/XGBoost tests (steps/save intervals).
- Data tests retagged from team:ml/ray_air to team:data.
Code cleanup:
- Remove construct_path from train/_internal/utils.py and its unit test; update pydoclint baseline accordingly.
Docs:
- Fix outdated example links and remove obsolete Tune example pages/links.
Bazel BUILD:
- Widespread tag/env updates; add/remove tests to reflect v2 enablement and GPU tagging changes.

^{Written by Cursor Bugbot for commit b88db74. This will update automatically on new commits. Configure here.}

Legacy V1 tests

This PR explicitly marked legacy v1 tests by running them with RAY_TRAIN_V2_ENABLED=0. See the BUILD files.

Followups

There are some tests where some tests should be migrated to V2 and some other tests in the file are legacy V1 tests. I've marked these as TODOs in the BUILD files to migrate in a follow-up.

Signed-off-by: Justin Yu <[email protected]>

gemini-code-assist

Code Review

This pull request introduces significant changes to enable Ray Train v2. The modifications include updating CI configurations, migrating numerous tests to use the v2 API by setting RAY_TRAIN_V2_ENABLED=1, and cleaning up obsolete v1 examples and tests. The changes are generally consistent and well-structured.

My main concern is the removal of tests in python/ray/train/tests/test_gpu_auto_transfer.py that verified critical GPU-related functionality. While the tests used v1 APIs, the underlying features are still important for v2. I've left a specific comment on this.

Otherwise, the migration strategy seems sound, with clear use of environment variables to toggle between API versions and TODOs marking areas for future work.

Signed-off-by: Justin Yu <[email protected]>

…n_enable_v2

Signed-off-by: Justin Yu <[email protected]>

#56868 replaced `gpu` tags with `train_v2_gpu` in order to separate out the tests into 2 CI pipelines. However, there's a py312 test that only runs in postmerge that filters out the `gpu` tag. This resulted in GPU tests running on the py312 CPU test runner. This PR fixes the issue by filtering out the new `train_v2_gpu` tag as well. Signed-off-by: Justin Yu <[email protected]>

Enable Train v2 across CI with new CPU/GPU jobs, migrate tests to v2 with env flags and tag updates, remove legacy Tune examples/utilities, and fix docs/BUILD configurations. --------- Signed-off-by: Justin Yu <[email protected]>

ray-project#56868 replaced `gpu` tags with `train_v2_gpu` in order to separate out the tests into 2 CI pipelines. However, there's a py312 test that only runs in postmerge that filters out the `gpu` tag. This resulted in GPU tests running on the py312 CPU test runner. This PR fixes the issue by filtering out the new `train_v2_gpu` tag as well. Signed-off-by: Justin Yu <[email protected]>

Enable Train v2 across CI with new CPU/GPU jobs, migrate tests to v2 with env flags and tag updates, remove legacy Tune examples/utilities, and fix docs/BUILD configurations. --------- Signed-off-by: Justin Yu <[email protected]> Signed-off-by: Josh Kodi <[email protected]>

ray-project#56868 replaced `gpu` tags with `train_v2_gpu` in order to separate out the tests into 2 CI pipelines. However, there's a py312 test that only runs in postmerge that filters out the `gpu` tag. This resulted in GPU tests running on the py312 CPU test runner. This PR fixes the issue by filtering out the new `train_v2_gpu` tag as well. Signed-off-by: Justin Yu <[email protected]> Signed-off-by: Josh Kodi <[email protected]>

Enable Train v2 across CI with new CPU/GPU jobs, migrate tests to v2 with env flags and tag updates, remove legacy Tune examples/utilities, and fix docs/BUILD configurations. --------- Signed-off-by: Justin Yu <[email protected]>

ray-project#56868 replaced `gpu` tags with `train_v2_gpu` in order to separate out the tests into 2 CI pipelines. However, there's a py312 test that only runs in postmerge that filters out the `gpu` tag. This resulted in GPU tests running on the py312 CPU test runner. This PR fixes the issue by filtering out the new `train_v2_gpu` tag as well. Signed-off-by: Justin Yu <[email protected]>

Ports over the remaining unit tests that were marked as TODOs from this series of PRs: #57534, #57256, #56868, #56820, #56816. Notably: * `test_new_dataset_config -> test_data_integration` * `test_backend -> test_torch_trainer, test_worker_group` * `test_gpu -> test_torch_gpu` This PR also finishes migrating the Tune LightGBM/Keras examples which were unblocked by #57042 and #57121. --------- Signed-off-by: Justin Yu <[email protected]>

Enable Train v2 across CI with new CPU/GPU jobs, migrate tests to v2 with env flags and tag updates, remove legacy Tune examples/utilities, and fix docs/BUILD configurations. --------- Signed-off-by: Justin Yu <[email protected]>

ray-project#56868 replaced `gpu` tags with `train_v2_gpu` in order to separate out the tests into 2 CI pipelines. However, there's a py312 test that only runs in postmerge that filters out the `gpu` tag. This resulted in GPU tests running on the py312 CPU test runner. This PR fixes the issue by filtering out the new `train_v2_gpu` tag as well. Signed-off-by: Justin Yu <[email protected]>

Ports over the remaining unit tests that were marked as TODOs from this series of PRs: ray-project#57534, ray-project#57256, ray-project#56868, ray-project#56820, ray-project#56816. Notably: * `test_new_dataset_config -> test_data_integration` * `test_backend -> test_torch_trainer, test_worker_group` * `test_gpu -> test_torch_gpu` This PR also finishes migrating the Tune LightGBM/Keras examples which were unblocked by ray-project#57042 and ray-project#57121. --------- Signed-off-by: Justin Yu <[email protected]>

Enable Train v2 across CI with new CPU/GPU jobs, migrate tests to v2 with env flags and tag updates, remove legacy Tune examples/utilities, and fix docs/BUILD configurations. --------- Signed-off-by: Justin Yu <[email protected]> Signed-off-by: xgui <[email protected]>

ray-project#56868 replaced `gpu` tags with `train_v2_gpu` in order to separate out the tests into 2 CI pipelines. However, there's a py312 test that only runs in postmerge that filters out the `gpu` tag. This resulted in GPU tests running on the py312 CPU test runner. This PR fixes the issue by filtering out the new `train_v2_gpu` tag as well. Signed-off-by: Justin Yu <[email protected]> Signed-off-by: xgui <[email protected]>

Ports over the remaining unit tests that were marked as TODOs from this series of PRs: ray-project#57534, ray-project#57256, ray-project#56868, ray-project#56820, ray-project#56816. Notably: * `test_new_dataset_config -> test_data_integration` * `test_backend -> test_torch_trainer, test_worker_group` * `test_gpu -> test_torch_gpu` This PR also finishes migrating the Tune LightGBM/Keras examples which were unblocked by ray-project#57042 and ray-project#57121. --------- Signed-off-by: Justin Yu <[email protected]> Signed-off-by: xgui <[email protected]>

Enable Train v2 across CI with new CPU/GPU jobs, migrate tests to v2 with env flags and tag updates, remove legacy Tune examples/utilities, and fix docs/BUILD configurations. --------- Signed-off-by: Justin Yu <[email protected]> Signed-off-by: elliot-barn <[email protected]>

#56868 replaced `gpu` tags with `train_v2_gpu` in order to separate out the tests into 2 CI pipelines. However, there's a py312 test that only runs in postmerge that filters out the `gpu` tag. This resulted in GPU tests running on the py312 CPU test runner. This PR fixes the issue by filtering out the new `train_v2_gpu` tag as well. Signed-off-by: Justin Yu <[email protected]> Signed-off-by: elliot-barn <[email protected]>

Ports over the remaining unit tests that were marked as TODOs from this series of PRs: #57534, #57256, #56868, #56820, #56816. Notably: * `test_new_dataset_config -> test_data_integration` * `test_backend -> test_torch_trainer, test_worker_group` * `test_gpu -> test_torch_gpu` This PR also finishes migrating the Tune LightGBM/Keras examples which were unblocked by #57042 and #57121. --------- Signed-off-by: Justin Yu <[email protected]> Signed-off-by: elliot-barn <[email protected]>

Enable Train v2 across CI with new CPU/GPU jobs, migrate tests to v2 with env flags and tag updates, remove legacy Tune examples/utilities, and fix docs/BUILD configurations. --------- Signed-off-by: Justin Yu <[email protected]>

ray-project#56868 replaced `gpu` tags with `train_v2_gpu` in order to separate out the tests into 2 CI pipelines. However, there's a py312 test that only runs in postmerge that filters out the `gpu` tag. This resulted in GPU tests running on the py312 CPU test runner. This PR fixes the issue by filtering out the new `train_v2_gpu` tag as well. Signed-off-by: Justin Yu <[email protected]>

Ports over the remaining unit tests that were marked as TODOs from this series of PRs: ray-project#57534, ray-project#57256, ray-project#56868, ray-project#56820, ray-project#56816. Notably: * `test_new_dataset_config -> test_data_integration` * `test_backend -> test_torch_trainer, test_worker_group` * `test_gpu -> test_torch_gpu` This PR also finishes migrating the Tune LightGBM/Keras examples which were unblocked by ray-project#57042 and ray-project#57121. --------- Signed-off-by: Justin Yu <[email protected]>

Enable Train v2 across CI with new CPU/GPU jobs, migrate tests to v2 with env flags and tag updates, remove legacy Tune examples/utilities, and fix docs/BUILD configurations. --------- Signed-off-by: Justin Yu <[email protected]> Signed-off-by: Aydin Abiar <[email protected]>

ray-project#56868 replaced `gpu` tags with `train_v2_gpu` in order to separate out the tests into 2 CI pipelines. However, there's a py312 test that only runs in postmerge that filters out the `gpu` tag. This resulted in GPU tests running on the py312 CPU test runner. This PR fixes the issue by filtering out the new `train_v2_gpu` tag as well. Signed-off-by: Justin Yu <[email protected]> Signed-off-by: Aydin Abiar <[email protected]>

Ports over the remaining unit tests that were marked as TODOs from this series of PRs: ray-project#57534, ray-project#57256, ray-project#56868, ray-project#56820, ray-project#56816. Notably: * `test_new_dataset_config -> test_data_integration` * `test_backend -> test_torch_trainer, test_worker_group` * `test_gpu -> test_torch_gpu` This PR also finishes migrating the Tune LightGBM/Keras examples which were unblocked by ray-project#57042 and ray-project#57121. --------- Signed-off-by: Justin Yu <[email protected]> Signed-off-by: Aydin Abiar <[email protected]>

Enable Train v2 across CI with new CPU/GPU jobs, migrate tests to v2 with env flags and tag updates, remove legacy Tune examples/utilities, and fix docs/BUILD configurations. --------- Signed-off-by: Justin Yu <[email protected]> Signed-off-by: Future-Outlier <[email protected]>

ray-project#56868 replaced `gpu` tags with `train_v2_gpu` in order to separate out the tests into 2 CI pipelines. However, there's a py312 test that only runs in postmerge that filters out the `gpu` tag. This resulted in GPU tests running on the py312 CPU test runner. This PR fixes the issue by filtering out the new `train_v2_gpu` tag as well. Signed-off-by: Justin Yu <[email protected]> Signed-off-by: Future-Outlier <[email protected]>

Ports over the remaining unit tests that were marked as TODOs from this series of PRs: ray-project#57534, ray-project#57256, ray-project#56868, ray-project#56820, ray-project#56816. Notably: * `test_new_dataset_config -> test_data_integration` * `test_backend -> test_torch_trainer, test_worker_group` * `test_gpu -> test_torch_gpu` This PR also finishes migrating the Tune LightGBM/Keras examples which were unblocked by ray-project#57042 and ray-project#57121. --------- Signed-off-by: Justin Yu <[email protected]> Signed-off-by: Future-Outlier <[email protected]>

justinvyu added 8 commits September 23, 2025 16:35

remove unnecessary tests

535afa5

Signed-off-by: Justin Yu <[email protected]>

first pass

3f2ae4e

Signed-off-by: Justin Yu <[email protected]>

2nd pass

9a0f75e

Signed-off-by: Justin Yu <[email protected]>

gpu_only???

8e996cd

Signed-off-by: Justin Yu <[email protected]>

update to v2 tags

f9577f5

Signed-off-by: Justin Yu <[email protected]>

pass 3

738f914

Signed-off-by: Justin Yu <[email protected]>

add train v2 gpu pipeline

eb2d040

Signed-off-by: Justin Yu <[email protected]>

remove legacy tune examples

4ad1e9e

Signed-off-by: Justin Yu <[email protected]>

justinvyu requested review from a team as code owners September 24, 2025 01:01

gemini-code-assist bot reviewed Sep 24, 2025

View reviewed changes

ray-gardener bot added the train Ray Train Related Issue label Sep 24, 2025

don't run wandb/comet on v2 pipeline

63c4ecd

Signed-off-by: Justin Yu <[email protected]>

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

justinvyu added 6 commits September 25, 2025 09:34

Merge branch 'master' of https://github.com/ray-project/ray into trai…

9edaee1

…n_enable_v2

add train_v2_gpu tag

b4c079f

Signed-off-by: Justin Yu <[email protected]>

fix ml rayci yml

14fe416

Signed-off-by: Justin Yu <[email protected]>

fix test_examples

d732ab5

Signed-off-by: Justin Yu <[email protected]>

fix test_gpu_examples

2c41fd8

Signed-off-by: Justin Yu <[email protected]>

update build

e69ca0f

Signed-off-by: Justin Yu <[email protected]>

justinvyu requested a review from a team as a code owner September 25, 2025 21:43

This comment was marked as outdated.

Sign in to view

justinvyu added 3 commits September 26, 2025 12:40

remove

8f9889e

Signed-off-by: Justin Yu <[email protected]>

fix hf test params

911abdd

Signed-off-by: Justin Yu <[email protected]>

remove ray.train.tensorflow.tensorflow_trainer imports

8666d4b

Signed-off-by: Justin Yu <[email protected]>

This comment was marked as outdated.

Sign in to view

no need to test tune x torch x train

e50bea6

Signed-off-by: Justin Yu <[email protected]>

This comment was marked as outdated.

Sign in to view

remove tune x tf x train

86d90d9

Signed-off-by: Justin Yu <[email protected]>

justinvyu merged commit f19664e into ray-project:master Oct 8, 2025
7 checks passed

justinvyu deleted the train_enable_v2 branch October 8, 2025 00:46

justinvyu mentioned this pull request Oct 8, 2025

[train] Fix py312 GPU test filter #57559

Merged

justinvyu mentioned this pull request Oct 15, 2025

[train] Port over remaining tests to Train v2 #57722

Merged

7 tasks

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[train] Enable v2 for ray/train/tests#56868

[train] Enable v2 for ray/train/tests#56868
justinvyu merged 30 commits intoray-project:masterfrom
justinvyu:train_enable_v2

justinvyu commented Sep 24, 2025 •

edited

Loading

Uh oh!

gemini-code-assist bot left a comment

Uh oh!

This comment was marked as outdated.

Uh oh!

This comment was marked as outdated.

Uh oh!

This comment was marked as outdated.

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Conversation

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

Uh oh!

Legacy V1 tests

Followups

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

This comment was marked as outdated.

Uh oh!

This comment was marked as outdated.

Uh oh!

This comment was marked as outdated.

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

justinvyu commented Sep 24, 2025 •

edited

Loading