[Tune][Air] Fix MLflowLoggerCallback to enable its use with PBT (#27783) by sjp611 · Pull Request #42182 · ray-project/ray

sjp611 · 2024-01-04T16:55:03Z

Why are these changes needed?

In the current MLflowLoggerCallback implementation, the function to store parameters in MLflow (mlflow_util.log_param) is only used in log_trial_start.

The problem is that PBT explores various parameters. The initially stored parameters are almost always different from the optimal parameters.

We need to be able to verify not only the initial parameters but also the optimal parameters in MLflow. The current implementation only allows us to verify the initial parameters.

This contribution modifies the MLflowLoggerCallback function to enable the storage of both the final and initial parameters in MLflow.

Related issue number

"Open #27783"

Checks

I've signed off every commit(by using the -s flag, i.e., git commit -s) in this PR.
I've run scripts/format.sh to lint the changes in this PR.
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 :(

…project#27783) Signed-off-by: Sungjoon Park <sjp611@gmail.com>

stale · 2024-04-22T12:00:40Z

This pull request has been automatically marked as stale because it has not had recent activity. It will be closed in 14 days if no further activity occurs. Thank you for your contributions.

If you'd like to keep this open, just leave any comment, and the stale label will be removed.

sjp611 · 2024-04-23T06:35:01Z

I have specifically modified my comment

github-actions · 2025-06-01T00:40:19Z

This pull request has been automatically marked as stale because it has not had
any activity for 14 days. It will be closed in another 14 days if no further activity occurs.
Thank you for your contributions.

You can always ask for help on our discussion forum or Ray's public slack channel.

If you'd like to keep this open, just leave any comment, and the stale label will be removed.

python/ray/air/integrations/mlflow.py

TimothySeah

Also please add a unit test for this, thanks!

…project#27783) Why are these changes needed? In the current MLflowLoggerCallback implementation, the function to store parameters in MLflow (mlflow_util.log_param) is only used in log_trial_start. The problem is that PBT explores various parameters. The initially stored parameters are almost always different from the optimal parameters. We need to be able to verify not only the initial parameters but also the optimal parameters in MLflow. The current implementation only allows us to verify the initial parameters. This contribution modifies the MLflowLoggerCallback function to enable the storage of both the final and initial parameters in MLflow. Related issue number "Open ray-project#27783" Checks - [x] I've signed off every commit(by using the -s flag, i.e., git commit -s) in this PR. - [x] I've run scripts/format.sh to lint the changes in this PR. - [ ] 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 - [x] Unit tests - [ ] Release tests - [ ] This PR is not tested :( Signed-off-by: Sung Joon Park <sjp611@gmail.com>

sjp611 · 2025-06-09T14:28:25Z

My apologies for the multiple notifications while I was cleaning up the PR 🥲

python/ray/air/tests/test_integration_mlflow.py

…project#27783) Why are these changes needed? In the current MLflowLoggerCallback implementation, the function to store parameters in MLflow (mlflow_util.log_param) is only used in log_trial_start. The problem is that PBT explores various parameters. The initially stored parameters are almost always different from the optimal parameters. We need to be able to verify not only the initial parameters but also the optimal parameters in MLflow. The current implementation only allows us to verify the initial parameters. This contribution modifies the MLflowLoggerCallback function to enable the storage of both the final and initial parameters in MLflow. Related issue number "Open ray-project#27783" Checks - [x] I've signed off every commit(by using the -s flag, i.e., git commit -s) in this PR. - [x] I've run scripts/format.sh to lint the changes in this PR. - [ ] 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 - [x] Unit tests - [ ] Release tests - [ ] This PR is not tested :( Signed-off-by: Sung Joon Park <sjp611@gmail.com>

…) (#42182) In the current MLflowLoggerCallback implementation, the function to store parameters in MLflow (mlflow_util.log_param) is only used in log_trial_start. The problem is that PBT explores various parameters. The initially stored parameters are almost always different from the optimal parameters. We need to be able to verify not only the initial parameters but also the optimal parameters in MLflow. The current implementation only allows us to verify the initial parameters. This contribution modifies the MLflowLoggerCallback function to enable the storage of both the final and initial parameters in MLflow. Signed-off-by: Sungjoon Park <sjp611@gmail.com> Signed-off-by: Sung Joon Park <sjp611@gmail.com> Signed-off-by: can <can@anyscale.com>

…) (#42182) In the current MLflowLoggerCallback implementation, the function to store parameters in MLflow (mlflow_util.log_param) is only used in log_trial_start. The problem is that PBT explores various parameters. The initially stored parameters are almost always different from the optimal parameters. We need to be able to verify not only the initial parameters but also the optimal parameters in MLflow. The current implementation only allows us to verify the initial parameters. This contribution modifies the MLflowLoggerCallback function to enable the storage of both the final and initial parameters in MLflow. Signed-off-by: Sungjoon Park <sjp611@gmail.com> Signed-off-by: Sung Joon Park <sjp611@gmail.com> Signed-off-by: elliot-barn <elliot.barnwell@anyscale.com>

sjp611 added 3 commits January 5, 2024 01:26

[Tune][Air] Fix MLflowLoggerCallback to enable its use with PBT (ray-…

ec13c83

…project#27783) Signed-off-by: Sungjoon Park <sjp611@gmail.com>

Merge branch 'master' into new-feature

75a1031

Merge branch 'master' into new-feature

4c51d2f

anyscalesam added the tune Tune-related issues label Feb 28, 2024

Merge branch 'master' into new-feature

d5afd03

Merge branch 'master' into new-feature

49e72ba

Merge branch 'master' into new-feature

db99487

anyscalesam added the triage Needs triage (eg: priority, bug/not-bug, and owning component) label May 14, 2024

Merge branch 'master' into new-feature

cf11ba8

hainesmichaelc added the community-contribution Contributed by the community label Apr 4, 2025

github-actions bot added the stale The issue is stale. It will be closed within 7 days unless there are further conversation label Jun 1, 2025

Merge branch 'master' into new-feature

13e84b0

github-actions bot removed the stale The issue is stale. It will be closed within 7 days unless there are further conversation label Jun 4, 2025

TimothySeah reviewed Jun 4, 2025

View reviewed changes

python/ray/air/integrations/mlflow.py Outdated Show resolved Hide resolved

TimothySeah reviewed Jun 6, 2025

View reviewed changes

Merge branch 'master' into new-feature

9ff9c87

sjp611 force-pushed the new-feature branch 2 times, most recently from 07e68b5 to 46e20ae Compare June 9, 2025 12:46

sjp611 requested review from a team, GeneDer, akshay-anyscale, edoakes, jjyao, richardliaw and zcin as code owners June 9, 2025 12:46

sjp611 requested a review from WangTaoTheTonic as a code owner June 9, 2025 12:46

sjp611 force-pushed the new-feature branch from 46e20ae to 07e68b5 Compare June 9, 2025 12:48

sjp611 closed this Jun 9, 2025

sjp611 deleted the new-feature branch June 9, 2025 12:50

sjp611 restored the new-feature branch June 9, 2025 12:54

sjp611 reopened this Jun 9, 2025

sjp611 force-pushed the new-feature branch from 07e68b5 to 9ff9c87 Compare June 9, 2025 12:57

sjp611 closed this Jun 9, 2025

sjp611 deleted the new-feature branch June 9, 2025 13:04

sjp611 restored the new-feature branch June 9, 2025 13:12

sjp611 reopened this Jun 9, 2025

eicherseiji removed request for a team June 9, 2025 14:48

sjp611 requested a review from TimothySeah June 9, 2025 23:13

TimothySeah approved these changes Jun 11, 2025

View reviewed changes

python/ray/air/tests/test_integration_mlflow.py Outdated Show resolved Hide resolved

sjp611 added 2 commits June 12, 2025 07:29

Merge branch 'master' into new-feature

8f63ccf

matthewdeng enabled auto-merge (squash) June 11, 2025 22:34

github-actions bot added the go add ONLY when ready to merge, run all tests label Jun 11, 2025

matthewdeng merged commit f0f29e9 into ray-project:master Jun 11, 2025
6 of 7 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[Tune][Air] Fix MLflowLoggerCallback to enable its use with PBT (#27783)#42182

[Tune][Air] Fix MLflowLoggerCallback to enable its use with PBT (#27783)#42182
matthewdeng merged 12 commits intoray-project:masterfrom
sjp611:new-feature

sjp611 commented Jan 4, 2024 •

edited

Loading

Uh oh!

stale bot commented Apr 22, 2024

Uh oh!

sjp611 commented Apr 23, 2024 •

edited

Loading

Uh oh!

github-actions bot commented Jun 1, 2025

Uh oh!

Uh oh!

TimothySeah left a comment

Uh oh!

sjp611 commented Jun 9, 2025 •

edited

Loading

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

Conversation

sjp611 commented Jan 4, 2024 • 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!

stale bot commented Apr 22, 2024

Uh oh!

sjp611 commented Apr 23, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

github-actions bot commented Jun 1, 2025

Uh oh!

Uh oh!

TimothySeah left a comment

Choose a reason for hiding this comment

Uh oh!

sjp611 commented Jun 9, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

sjp611 commented Jan 4, 2024 •

edited

Loading

sjp611 commented Apr 23, 2024 •

edited

Loading

sjp611 commented Jun 9, 2025 •

edited

Loading