setup_mlflow API Ray Train V2 Compatibility

[JasonLi1909]

Users have run into issues using the setup_mlflow API with Ray Train V2 enabled:

• When using setup_mlflow in a tune run, a DeprecationWarning is raised because ray.train.get_context() is called from with a Ray Tune run
• When using setup_mlflow in a regular Train run, a DeprecationWarning is raised because TrainContext.get_trial_id() is called but this is no longer supported in Train V2

This PR:

• Makes setup_mlflow compatible with Ray Train V2 by avoiding deprecated API usage
• Adds tests for testing the compatibility of Train V2 with setup_mlflow

[gemini-code-assist]

Code Review

This pull request aims to add compatibility with Ray Train V2 for the setup_mlflow API by handling deprecation warnings. The changes introduce try-except blocks to catch DeprecationWarning and fall back to alternative APIs or gracefully handle the situation.

While the intent is good, I've found a critical issue where the new structure could lead to a NameError if ray.train.get_context() raises a RuntimeError. I've provided a suggestion to restructure the code to prevent this.

Additionally, I've pointed out that simply catching DeprecationWarning for context methods like get_world_rank() and then proceeding without trial information could lead to silent failures in MLflow logging. It would be better to use the newer APIs if possible.

[justinvyu]

Thanks, just a few minor things

[justinvyu]

Thanks!