[Tune]: MLflowLoggerCallback does not use tracking_uri

[hahahannes]

What happened + What you expected to happen

What happened: I am trying to use the MLflowLoggerCallback to log a figure during Tune trials. I noticed, that the artifacts did not get saved in the run data within MLFLOW. Therefore I checked whether the callback uses the correct MLFLOW tracking URI. Indeed it uses a file path like file:///..../ray_results/run/run_e82f3_00000_0_2023-06-09_16-14-12/mlruns.

What you expected to happen: The tracking URI should have been the one that I provided in the callback options.

Versions / Dependencies

ray==2.5.0 mlflow==2.4.0

python: Python 3.8.16 OS: Pop!_OS 22.04 LTS

Reproduction script

import mlflow
from ray import air, tune
from ray.air.integrations.mlflow import MLflowLoggerCallback, setup_mlflow

def run(config):
    print(f"MLFLOW URI: {mlflow.get_tracking_uri()}")

MLFLOW_URL = "http://localhost:5003"

tuner = tune.Tuner(
    tune.with_parameters(run),
    param_space={},
    run_config=air.RunConfig(
        name="run",
        verbose=2,
        callbacks=[MLflowLoggerCallback(
            tracking_uri=MLFLOW_URL,
            experiment_name="Run",
            save_artifact=True,
        )]
    )
)
tuner.fit()

When you run the script, you will expect the output MLFLOW URI: http://localhost:5003.

This works when the setup_mlflow function is used like:

def run_with_setup(config, mlflow_config):
    setup_mlflow(**mlflow_config)
    print(f"MLFLOW URI: {mlflow.get_tracking_uri()}")


mlflow_config = {"experiment_name": "Run", "tracking_uri": MLFLOW_URL}
tuner = tune.Tuner(
    tune.with_parameters(run_with_setup, mlflow_config=mlflow_config),
    param_space={},
    run_config=air.RunConfig(name="run", verbose=2)
)
tuner.fit()

Issue Severity

Low: It annoys or frustrates me.

[hahahannes]

I am also happy to contribute in case this is a bug :)

[krfricke]

Hi @hahahannes,

the MLFLowLoggerCallback runs on the driver (where your tune script is running). It does not setup a session on the trainable. That's precisely what setup_mlflow is for.

In a nutshell: The logger callback automatically reports result and checkpoints. If this is all you need, you don't need to access the mlflow session in the trainable at all.

If you need more fine grained control, you shouldn't use the logger callback, but use e.g. setup_mlflow instead.

I think we can make this distinction more clear in the docs. If you'd like to contribute that, it would be much appreciated :-)

[hahahannes]

I see! Thank you very much for the clarification. Just for completeness: I like the callback as it does most of the logging automatically. But I wanted to log a figure additionally with a mlflow session and thought I could reuse the Callback.

Will open a Pull Request!

[krfricke]

Thanks! Please feel free to tag me on the PR. I'm closing this issue for now