[train][tune] Refactor and Reintroduce keras tune callback by liulehui · Pull Request #57121 · ray-project/ray

liulehui · 2025-10-02T03:18:12Z

This PR re-introduces the ray.tune.integration.keras import TuneReportCheckpointCallback
This is mainly uses for Tune only usage, testing through script here: https://gist.github.com/liulehui/99a5560031c10c274c1d7d99247033d9

Why are these changes needed?

Related issue number

Checks

I've signed off every commit(by using the -s flag, i.e., git commit -s) in this PR.
I've run pre-commit jobs to lint the changes in this PR. (pre-commit setup)
I've included any doc changes needed for https://docs.ray.io/en/master/.
- I've added any new APIs to the API Reference. For example, if I added a
  method in Tune, I've added it in doc/source/tune/api/ under the
  corresponding .rst file.
I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/
Testing Strategy
- Unit tests
- Release tests
- This PR is not tested :(

justinvyu

Thanks!

python/ray/air/integrations/keras.py

python/ray/tune/integration/keras.py

Signed-off-by: Lehui Liu <lehui@anyscale.com>

cursor · 2025-10-06T16:23:15Z

python/ray/tune/integration/keras.py

+    def _save_and_report_checkpoint(
+        self, metrics: Dict, checkpoint: TensorflowCheckpoint
+    ):
+        ray.tune.report(metrics, checkpoint=checkpoint)


Bug: Tune Callback Passes Invalid Argument

The TuneReportCheckpointCallback passes a checkpoint argument to ray.tune.report(). This causes a TypeError because ray.tune.report() does not accept a checkpoint parameter, unlike ray.train.report(). Tune's API handles checkpoints through a different mechanism.

justinvyu

Thanks!

…ct#57121) 1. This PR re-introduces the `ray.tune.integration.keras import TuneReportCheckpointCallback` 2. This is mainly uses for Tune only usage, testing through script here: https://gist.github.com/liulehui/99a5560031c10c274c1d7d99247033d9 --------- Signed-off-by: Lehui Liu <lehui@anyscale.com>

…ct#57121) 1. This PR re-introduces the `ray.tune.integration.keras import TuneReportCheckpointCallback` 2. This is mainly uses for Tune only usage, testing through script here: https://gist.github.com/liulehui/99a5560031c10c274c1d7d99247033d9 --------- Signed-off-by: Lehui Liu <lehui@anyscale.com> Signed-off-by: Josh Kodi <joshkodi@gmail.com>

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 <justinvyu@anyscale.com>

…ct#57121) 1. This PR re-introduces the `ray.tune.integration.keras import TuneReportCheckpointCallback` 2. This is mainly uses for Tune only usage, testing through script here: https://gist.github.com/liulehui/99a5560031c10c274c1d7d99247033d9 --------- Signed-off-by: Lehui Liu <lehui@anyscale.com>

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 <justinvyu@anyscale.com>

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 <justinvyu@anyscale.com> Signed-off-by: xgui <xgui@anyscale.com>

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 <justinvyu@anyscale.com> Signed-off-by: elliot-barn <elliot.barnwell@anyscale.com>

…ct#57121) 1. This PR re-introduces the `ray.tune.integration.keras import TuneReportCheckpointCallback` 2. This is mainly uses for Tune only usage, testing through script here: https://gist.github.com/liulehui/99a5560031c10c274c1d7d99247033d9 --------- Signed-off-by: Lehui Liu <lehui@anyscale.com>

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 <justinvyu@anyscale.com>

…ct#57121) 1. This PR re-introduces the `ray.tune.integration.keras import TuneReportCheckpointCallback` 2. This is mainly uses for Tune only usage, testing through script here: https://gist.github.com/liulehui/99a5560031c10c274c1d7d99247033d9 --------- Signed-off-by: Lehui Liu <lehui@anyscale.com> Signed-off-by: Aydin Abiar <aydin@anyscale.com>

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 <justinvyu@anyscale.com> Signed-off-by: Aydin Abiar <aydin@anyscale.com>

…ct#57121) 1. This PR re-introduces the `ray.tune.integration.keras import TuneReportCheckpointCallback` 2. This is mainly uses for Tune only usage, testing through script here: https://gist.github.com/liulehui/99a5560031c10c274c1d7d99247033d9 --------- Signed-off-by: Lehui Liu <lehui@anyscale.com> Signed-off-by: Future-Outlier <eric901201@gmail.com>

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 <justinvyu@anyscale.com> Signed-off-by: Future-Outlier <eric901201@gmail.com>

liulehui changed the title ~~[train][tune] refactor and reintroduce keras tune callback~~ [train][tune] Refactor and Reintroduce keras tune callback Oct 2, 2025

liulehui marked this pull request as ready for review October 2, 2025 17:35

liulehui requested review from a team as code owners October 2, 2025 17:35

ray-gardener bot added the tune Tune-related issues label Oct 2, 2025

justinvyu reviewed Oct 3, 2025

View reviewed changes

python/ray/air/integrations/keras.py Show resolved Hide resolved

python/ray/tune/integration/keras.py Outdated Show resolved Hide resolved

liulehui added the go add ONLY when ready to merge, run all tests label Oct 6, 2025

This comment was marked as outdated.

Sign in to view

liulehui added 4 commits October 6, 2025 09:04

refactor and reintroduce keras tune callback

a9959e1

Signed-off-by: Lehui Liu <lehui@anyscale.com>

formatting

2c7ef60

Signed-off-by: Lehui Liu <lehui@anyscale.com>

not using air

af0d8ed

Signed-off-by: Lehui Liu <lehui@anyscale.com>

address comments

834726e

Signed-off-by: Lehui Liu <lehui@anyscale.com>

liulehui force-pushed the fix-keras-callback branch from 9dff146 to 834726e Compare October 6, 2025 16:04

fix lint

f66fc50

Signed-off-by: Lehui Liu <lehui@anyscale.com>

This comment was marked as outdated.

Sign in to view

address comment

643ca06

Signed-off-by: Lehui Liu <lehui@anyscale.com>

cursor bot reviewed Oct 6, 2025

View reviewed changes

justinvyu approved these changes Oct 6, 2025

View reviewed changes

justinvyu merged commit f6a83c0 into ray-project:master Oct 6, 2025
6 checks passed

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][tune] Refactor and Reintroduce keras tune callback#57121

[train][tune] Refactor and Reintroduce keras tune callback#57121
justinvyu merged 6 commits intoray-project:masterfrom
liulehui:fix-keras-callback

liulehui commented Oct 2, 2025 •

edited

Loading

Uh oh!

justinvyu left a comment

Uh oh!

Uh oh!

Uh oh!

This comment was marked as outdated.

Uh oh!

This comment was marked as outdated.

Uh oh!

cursor bot Oct 6, 2025

Uh oh!

justinvyu left a comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

liulehui commented Oct 2, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Why are these changes needed?

Related issue number

Checks

Uh oh!

justinvyu left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

This comment was marked as outdated.

Uh oh!

This comment was marked as outdated.

Uh oh!

cursor bot Oct 6, 2025

Choose a reason for hiding this comment

Bug: Tune Callback Passes Invalid Argument

Uh oh!

justinvyu left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

liulehui commented Oct 2, 2025 •

edited

Loading