Repository: optuna/optuna

Issue:
MLFlow logging multiple metrics
I want to log multiple metrics in MLFlow but under the hood, Optuna uses mlflow.log_metric in its MLFlowCallBack so there is no straigthforward way to do it.

## Motivation

When I am using MLFlow to log experiments, I might want to log additional metrics even though I am not optimizing specifically for them. Or I could be using their sum in my objective function and I might want to learn their individual contribution to the objective score.

## Description

MLFlowCallBack can get additional parameters (functional parameters, not hyperparameters) to log.

## Alternatives (optional)

--

## Additional context (optional)

--



Discussion:
Thanks for the proposal. The proposed feature is reasonable to me. Let me label this issue as "contribution-welcome".
This issue has not seen any recent activity.
Hi, I run into the same case. I am thinking about writing a decorator. If it works I voluntary share the result with you.
@HideakiImamura A decorator for the objective function seems the work fine and address the problem. Should I try to work out a PR or is totally out of scope where you want to go with optuna?
@lucafurrer Thank you for the investigation! Could you create a draft PR? I'd like to discuss the design in the draft PR.
Hi @HideakiImamura I will do that (I hope soon 😉)

Resolving patch:
diff --git a/optuna/integration/mlflow.py b/optuna/integration/mlflow.py
--- a/optuna/integration/mlflow.py
+++ b/optuna/integration/mlflow.py
@@ -1,17 +1,23 @@
 import textwrap
+from typing import Callable
 from typing import Dict
 from typing import Optional
+from typing import Sequence
+from typing import Union
 
 import optuna
 from optuna._experimental import experimental
 from optuna._imports import try_import
 from optuna._study_direction import StudyDirection
+from optuna.study import ObjectiveFuncType
 from optuna.trial import TrialState
 
 
 with try_import() as _imports:
     import mlflow
 
+RUN_ID_ATTRIBUTE_KEY = "mlflow_run_id"
+
 
 @experimental("1.4.0")
 class MLflowCallback(object):
@@ -64,6 +70,39 @@ def objective(trial):
 
             INFO: 'my_study' does not exist. Creating a new experiment
 
+        Add additional logging to MLflow
+
+        .. testcode::
+
+            import optuna
+            import mlflow
+            from optuna.integration.mlflow import MLflowCallback
+
+            mlflc = MLflowCallback(
+                tracking_uri=YOUR_TRACKING_URI,
+                metric_name="my metric score",
+            )
+
+
+            @mlflc.track_in_mlflow()
+            def objective(trial):
+                x = trial.suggest_float("x", -10, 10)
+                mlflow.log_param("power", 2)
+                mlflow.log_metric("base of metric", x - 2)
+
+                return (x - 2) ** 2
+
+
+            study = optuna.create_study(study_name="my_other_study")
+            study.optimize(objective, n_trials=10, callbacks=[mlflc])
+
+
+        .. testoutput::
+            :hide:
+            :options: +NORMALIZE_WHITESPACE
+
+            INFO: 'my_other_study' does not exist. Creating a new experiment
+
     Args:
         tracking_uri:
             The URI of the MLflow tracking server.
@@ -88,7 +127,7 @@ def objective(trial):
     def __init__(
         self,
         tracking_uri: Optional[str] = None,
-        metric_name: str = "value",
+        metric_name: Union[str] = "value",
         nest_trials: bool = False,
         tag_study_user_attrs: bool = False,
     ) -> None:
@@ -102,56 +141,122 @@ def __init__(
 
     def __call__(self, study: optuna.study.Study, trial: optuna.trial.FrozenTrial) -> None:
 
-        # This sets the tracking_uri for MLflow.
-        if self._tracking_uri is not None:
-            mlflow.set_tracking_uri(self._tracking_uri)
-
-        # This sets the experiment of MLflow.
-        mlflow.set_experiment(study.study_name)
+        self._initialize_experiment(study)
 
-        with mlflow.start_run(run_name=str(trial.number), nested=self._nest_trials):
+        with mlflow.start_run(
+            run_id=trial.system_attrs.get(RUN_ID_ATTRIBUTE_KEY),
+            run_name=str(trial.number),
+            nested=self._nest_trials,
+        ):
 
             # This sets the metric for MLflow.
-            trial_value = trial.value if trial.value is not None else float("nan")
-            mlflow.log_metric(self._metric_name, trial_value)
+            self._log_metric(trial.value)
 
             # This sets the params for MLflow.
-            mlflow.log_params(trial.params)
+            self._log_params(trial)
 
             # This sets the tags for MLflow.
-            tags: Dict[str, str] = {}
-            tags["number"] = str(trial.number)
-            tags["datetime_start"] = str(trial.datetime_start)
-            tags["datetime_complete"] = str(trial.datetime_complete)
-
-            # Set state and convert it to str and remove the common prefix.
-            trial_state = trial.state
-            if isinstance(trial_state, TrialState):
-                tags["state"] = str(trial_state).split(".")[-1]
-
-            # Set direction and convert it to str and remove the common prefix.
-            study_direction = study.direction
-            if isinstance(study_direction, StudyDirection):
-                tags["direction"] = str(study_direction).split(".")[-1]
-
-            tags.update(trial.user_attrs)
-            distributions = {
-                (k + "_distribution"): str(v) for (k, v) in trial.distributions.items()
-            }
-            tags.update(distributions)
-
-            if self._tag_study_user_attrs:
-                tags.update(study.user_attrs)
-
-            # This is a temporary fix on Optuna side. It avoids an error with user
-            # attributes that are too long. It should be fixed on MLflow side later.
-            # When it is fixed on MLflow side this codeblock can be removed.
-            # see https://github.com/optuna/optuna/issues/1340
-            # see https://github.com/mlflow/mlflow/issues/2931
-            max_mlflow_tag_length = 5000
-            for key, value in tags.items():
-                value = str(value)  # make sure it is a string
-                if len(value) > max_mlflow_tag_length:
-                    tags[key] = textwrap.shorten(value, max_mlflow_tag_length)
-
-            mlflow.set_tags(tags)
+            self._set_tags(trial, study)
+
+    @experimental("2.9.0")
+    def track_in_mlflow(self) -> Callable:
+        """Decorator for using MLFlow logging in the objective function.
+
+        This decorator enables the extension of MLFlow logging provided by the callback.
+
+        All information logged in the decorated objective function will be added to the MLFlow
+        run for the trial created by the callback.
+
+        Returns:
+            ObjectiveFuncType: Objective function with tracking to MLFlow enabled
+        """
+
+        def decorator(func: ObjectiveFuncType) -> ObjectiveFuncType:
+            def wrapper(trial: optuna.trial.Trial) -> Union[float, Sequence[float]]:
+                study = trial.study
+                self._initialize_experiment(study)
+
+                with mlflow.start_run(run_name=str(trial.number), nested=self._nest_trials) as run:
+                    trial.set_system_attr(RUN_ID_ATTRIBUTE_KEY, run.info.run_id)
+
+                    return func(trial)
+
+            return wrapper
+
+        return decorator
+
+    def _initialize_experiment(self, study: optuna.study.Study) -> None:
+        """Initialize an MLFlow experiment with the study name.
+
+        If a tracking uri has been provided, MLFlow will be initialized to use it.
+
+        Args:
+            study: Study to be tracked in MLFlow
+        """
+        # This sets the tracking_uri for MLflow.
+        if self._tracking_uri is not None:
+            mlflow.set_tracking_uri(self._tracking_uri)
+
+        # This sets the experiment of MLflow.
+        mlflow.set_experiment(study.study_name)
+
+    def _set_tags(self, trial: optuna.trial.FrozenTrial, study: optuna.study.Study) -> None:
+        """Sets the Optuna tags for the current MLFlow run
+
+        Args:
+            trial: Trial to be tracked
+            study: Study to be tracked
+        """
+        tags: Dict[str, str] = {}
+        tags["number"] = str(trial.number)
+        tags["datetime_start"] = str(trial.datetime_start)
+
+        tags["datetime_complete"] = str(trial.datetime_complete)
+
+        # Set state and convert it to str and remove the common prefix.
+        trial_state = trial.state
+        if isinstance(trial_state, TrialState) and trial_state.is_finished():
+            tags["state"] = str(trial_state).split(".")[-1]
+
+        # Set direction and convert it to str and remove the common prefix.
+        study_direction = study.direction
+        if isinstance(study_direction, StudyDirection):
+            tags["direction"] = str(study_direction).split(".")[-1]
+
+        tags.update(trial.user_attrs)
+        distributions = {(k + "_distribution"): str(v) for (k, v) in trial.distributions.items()}
+        tags.update(distributions)
+
+        if self._tag_study_user_attrs:
+            tags.update(study.user_attrs)
+
+        # This is a temporary fix on Optuna side. It avoids an error with user
+        # attributes that are too long. It should be fixed on MLflow side later.
+        # When it is fixed on MLflow side this codeblock can be removed.
+        # see https://github.com/optuna/optuna/issues/1340
+        # see https://github.com/mlflow/mlflow/issues/2931
+        for key, value in tags.items():
+            value = str(value)  # make sure it is a string
+            if len(value) > mlflow.utils.validation.MAX_TAG_VAL_LENGTH:
+                tags[key] = textwrap.shorten(value, mlflow.utils.validation.MAX_TAG_VAL_LENGTH)
+
+        # This sets the tags for MLflow.
+        mlflow.set_tags(tags)
+
+    def _log_metric(self, value: Union[float, None]) -> None:
+        """Log the trial result as metric to MLFlow
+
+        Args:
+            value: Result of trial
+        """
+        trial_value = value if value is not None else float("nan")
+        mlflow.log_metric(self._metric_name, trial_value)
+
+    @staticmethod
+    def _log_params(trial: optuna.trial.BaseTrial) -> None:
+        """Log the parameters of the trial to MLFlow
+
+        Args:
+            trial: Trial to be tracked
+        """
+        mlflow.log_params(trial.params)


Test patch:
diff --git a/tests/integration_tests/test_mlflow.py b/tests/integration_tests/test_mlflow.py
--- a/tests/integration_tests/test_mlflow.py
+++ b/tests/integration_tests/test_mlflow.py
@@ -1,6 +1,7 @@
 import mlflow
 from mlflow.tracking import MlflowClient
 from mlflow.utils.mlflow_tags import MLFLOW_PARENT_RUN_ID
+import numpy as np
 import py
 import pytest
 
@@ -210,3 +211,169 @@ def test_tag_study_user_attrs(tmpdir: py.path.local, tag_study_user_attrs: bool)
         assert all((r.data.tags["my_study_attr"] == "a") for r in runs)
     else:
         assert all(("my_study_attr" not in r.data.tags) for r in runs)
+
+
+def test_track_in_mlflow_decorator(tmpdir: py.path.local) -> None:
+    tracking_file_name = "file:{}".format(tmpdir)
+    study_name = "my_study"
+    n_trials = 3
+
+    metric_name = "additional_metric"
+    metric = 3.14
+
+    mlflc = MLflowCallback(tracking_uri=tracking_file_name)
+
+    @mlflc.track_in_mlflow()
+    def _objective_func(trial: optuna.trial.Trial) -> float:
+
+        x = trial.suggest_float("x", -1.0, 1.0)
+        y = trial.suggest_float("y", 20, 30, log=True)
+        z = trial.suggest_categorical("z", (-1.0, 1.0))
+        assert isinstance(z, float)
+        trial.set_user_attr("my_user_attr", "my_user_attr_value")
+        mlflow.log_metric(metric_name, metric)
+        return (x - 2) ** 2 + (y - 25) ** 2 + z
+
+    study = optuna.create_study(study_name=study_name)
+    study.optimize(_objective_func, n_trials=n_trials, callbacks=[mlflc])
+
+    mlfl_client = MlflowClient(tracking_file_name)
+    experiments = mlfl_client.list_experiments()
+    assert len(experiments) == 1
+
+    experiment = experiments[0]
+    assert experiment.name == study_name
+    experiment_id = experiment.experiment_id
+
+    run_infos = mlfl_client.list_run_infos(experiment_id)
+    assert len(run_infos) == n_trials
+
+    first_run_id = run_infos[0].run_id
+    first_run = mlfl_client.get_run(first_run_id)
+    first_run_dict = first_run.to_dictionary()
+
+    assert metric_name in first_run_dict["data"]["metrics"]
+    assert first_run_dict["data"]["metrics"][metric_name] == metric
+
+
+def test_initialize_experiment(tmpdir: py.path.local) -> None:
+    tracking_file_name = "file:{}".format(tmpdir)
+    metric_name = "my_metric_name"
+    study_name = "my_study"
+
+    mlflc = MLflowCallback(tracking_uri=tracking_file_name, metric_name=metric_name)
+    study = optuna.create_study(study_name=study_name)
+
+    mlflc._initialize_experiment(study)
+
+    mlfl_client = MlflowClient(tracking_file_name)
+    experiments = mlfl_client.list_experiments()
+    assert len(experiments) == 1
+
+    experiment = experiments[0]
+    assert experiment.name == study_name
+
+
+def test_log_metric(tmpdir: py.path.local) -> None:
+    tracking_file_name = "file:{}".format(tmpdir)
+    metric_name = "my_metric_name"
+    study_name = "my_study"
+    metric_value = 3.17
+
+    mlflc = MLflowCallback(tracking_uri=tracking_file_name, metric_name=metric_name)
+    study = optuna.create_study(study_name=study_name)
+    mlflc._initialize_experiment(study)
+
+    with mlflow.start_run():
+        mlflc._log_metric(metric_value)
+
+    mlfl_client = MlflowClient(tracking_file_name)
+    experiments = mlfl_client.list_experiments()
+    experiment = experiments[0]
+    experiment_id = experiment.experiment_id
+
+    run_infos = mlfl_client.list_run_infos(experiment_id)
+    assert len(run_infos) == 1
+
+    first_run_id = run_infos[0].run_id
+    first_run = mlfl_client.get_run(first_run_id)
+    first_run_dict = first_run.to_dictionary()
+
+    assert metric_name in first_run_dict["data"]["metrics"]
+    assert first_run_dict["data"]["metrics"][metric_name] == metric_value
+
+
+def test_log_metric_none(tmpdir: py.path.local) -> None:
+    tracking_file_name = "file:{}".format(tmpdir)
+    metric_name = "my_metric_name"
+    study_name = "my_study"
+    metric_value = None
+
+    mlflc = MLflowCallback(tracking_uri=tracking_file_name, metric_name=metric_name)
+    study = optuna.create_study(study_name=study_name)
+    mlflc._initialize_experiment(study)
+
+    with mlflow.start_run():
+        mlflc._log_metric(metric_value)
+
+    mlfl_client = MlflowClient(tracking_file_name)
+    experiments = mlfl_client.list_experiments()
+    experiment = experiments[0]
+    experiment_id = experiment.experiment_id
+
+    run_infos = mlfl_client.list_run_infos(experiment_id)
+    assert len(run_infos) == 1
+
+    first_run_id = run_infos[0].run_id
+    first_run = mlfl_client.get_run(first_run_id)
+    first_run_dict = first_run.to_dictionary()
+
+    assert metric_name in first_run_dict["data"]["metrics"]
+    assert np.isnan(first_run_dict["data"]["metrics"][metric_name])
+
+
+def test_log_params(tmpdir: py.path.local) -> None:
+    tracking_file_name = "file:{}".format(tmpdir)
+    metric_name = "my_metric_name"
+    study_name = "my_study"
+
+    param1_name = "my_param1"
+    param1_value = "a"
+    param2_name = "my_param2"
+    param2_value = 5
+
+    params = {param1_name: param1_value, param2_name: param2_value}
+
+    mlflc = MLflowCallback(tracking_uri=tracking_file_name, metric_name=metric_name)
+    study = optuna.create_study(study_name=study_name)
+    mlflc._initialize_experiment(study)
+
+    with mlflow.start_run():
+
+        trial = optuna.trial.create_trial(
+            params=params,
+            distributions={
+                param1_name: optuna.distributions.CategoricalDistribution(["a", "b"]),
+                param2_name: optuna.distributions.UniformDistribution(0, 10),
+            },
+            value=5.0,
+        )
+        mlflc._log_params(trial)
+
+    mlfl_client = MlflowClient(tracking_file_name)
+    experiments = mlfl_client.list_experiments()
+    experiment = experiments[0]
+    experiment_id = experiment.experiment_id
+
+    run_infos = mlfl_client.list_run_infos(experiment_id)
+    assert len(run_infos) == 1
+
+    first_run_id = run_infos[0].run_id
+    first_run = mlfl_client.get_run(first_run_id)
+    first_run_dict = first_run.to_dictionary()
+
+    assert param1_name in first_run_dict["data"]["params"]
+    assert first_run_dict["data"]["params"][param1_name] == param1_value
+
+    assert param2_name in first_run_dict["data"]["params"]
+    assert first_run_dict["data"]["params"][param2_name] == str(param2_value)