Ray Tune train.report logs parameters via MLflowLoggerCallback in mlflow as metrics

Ray Libraries (Data, Train, Tune, Serve) Ray Tune
1

Hi,

I have a question to the functionality of ray tune in combination with mlflow.
Via the MLflowLoggerCallback API I want to start ray tune runs and save the parameters as well as the metrics to mlflow. I have a train_function and a tune_with_callback_function (see below). Within the train_function I use train.report to log the metrics to mlflow. But as the result ray tune not only logs the metric I passed to the train.report but also the params within the metrics in mlflow (see screenshot). Isn´t it overhead to pass the params to mlflow “Parameters” and also to “Metrics”? Within the Metrics I prefer to see only the wished metrics.
Is this a bug or maybe a feature - and is there a workaround to avoid the parameters within the Metrics?

I also tried to remove the train.report but then I only get the params and no metrics.

Bild_2024-01-27_103523955
Bild_2024-01-27_103523955392×595 9.57 KB

My specifications:

  • mlflow 2.9.2
  • ray 2.9.1

As example I used the ray.tune tutorial

def train_function(config):
    width, height = config["width"], config["height"]
    print(config)

    train.report({"my_metrics": 555})

    for step in range(config.get("steps", 100)):
        # Iterative training function - can be any arbitrary training procedure
        intermediate_score = evaluation_fn(step, width, height)
        # Feed the score back to Tune.
        train.report({"iterations": step, "mean_loss": intermediate_score})
        time.sleep(0.1)


def tune_with_callback(mlflow_tracking_uri, finish_fast=False):
    tuner = tune.Tuner(
        train_function,
        tune_config=tune.TuneConfig(num_samples=5),
        run_config=train.RunConfig(
            name="mlflow",
            callbacks=[
                MLflowLoggerCallback(
                    tracking_uri=mlflow_tracking_uri,
                    experiment_name="Tutorial_Callback_Tuner",
                    save_artifact=True,
                )
            ],
        ),
        param_space={
            "width": tune.randint(10, 100),
            "height": tune.randint(0, 100),
            "steps": 5 if finish_fast else 100,
        },
    )
    results = tuner.fit()

Thanks if someone could help!