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

[sjp611]

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 :(

[stale]

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]

I have specifically modified my comment