Repository: optuna/optuna
Pull request: MLFlow decorator for optimization function
<!-- Thank you for creating a pull request! In general, we merge your pull requests after they get two or more approvals. -->

## Motivation
With the callback integration of MLFlow it is not possible to add additional parameter, tags and metrics. Quite often you want to keep track of more information that actually available to Optuna. P.ex
- Additional metric which is not optimized
- Static parameters
- Evolution of metric over iterations in the algorithm
- Figures

fixes #2567 

## Description of the changes
- [x] Add a function usable as decorator to integration/mlflow
- [x] Restructured current code to avoid duplication
- [x] Unit-tests
- [x] Complete documentation

Problem statement:
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)

--

Hints from issue 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 😉)
Hi @HideakiImamura  I finally managed to make a draft PR #2670 with my idea. I would be glad if you could have a look at it.

Commit message: fix annotations

Patch to review:
diff --git a/optuna/integration/mlflow.py b/optuna/integration/mlflow.py
index 3412059e13e..9c664c81fb3 100644
--- a/optuna/integration/mlflow.py
+++ b/optuna/integration/mlflow.py
@@ -1,4 +1,5 @@
 import textwrap
+from typing import Callable
 from typing import Dict
 from typing import Optional
 
@@ -13,8 +14,79 @@
     import mlflow
 
 
+class MLFlowIntegrator(object):
+    def __init__(
+        self,
+        tracking_uri: Optional[str] = None,
+        metric_name: str = "value",
+        nest_trials: bool = False,
+        tag_study_user_attrs: bool = False,
+    ) -> None:
+
+        _imports.check()
+
+        self._tracking_uri = tracking_uri
+        self._metric_name = metric_name
+        self._nest_trials = nest_trials
+        self._tag_study_user_attrs = tag_study_user_attrs
+
+    def set_experiment(self, study: optuna.study.Study) -> 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)
+
+    def calculate_tags(self, trial: optuna.trial.BaseTrial, study: optuna.study.Study) -> None:
+
+        tags: Dict[str, str] = {}
+        tags["number"] = str(trial.number)
+        tags["datetime_start"] = str(trial.datetime_start)
+        if isinstance(trial, optuna.trial.FrozenTrial):
+            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
+        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)
+
+        # This sets the tags for MLflow.
+        mlflow.set_tags(tags)
+
+    def log_metric(self, value: float) -> None:
+        trial_value = value if value is not None else float("nan")
+        mlflow.log_metric(self._metric_name, trial_value)
+
+    def log_params(self, trial: optuna.trial.BaseTrial) -> None:
+        mlflow.log_params(trial.params)
+
+
 @experimental("1.4.0")
-class MLflowCallback(object):
+class MLflowCallback(MLFlowIntegrator):
     """Callback to track Optuna trials with MLflow.
 
     This callback adds relevant information that is
@@ -85,73 +157,42 @@ def objective(trial):
             set, key value pairs in study.user_attrs will supersede existing tags.
     """
 
-    def __init__(
-        self,
-        tracking_uri: Optional[str] = None,
-        metric_name: str = "value",
-        nest_trials: bool = False,
-        tag_study_user_attrs: bool = False,
-    ) -> None:
-
-        _imports.check()
-
-        self._tracking_uri = tracking_uri
-        self._metric_name = metric_name
-        self._nest_trials = nest_trials
-        self._tag_study_user_attrs = tag_study_user_attrs
-
     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.set_experiment(study)
 
         with mlflow.start_run(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)
+            if 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)
+            self.calculate_tags(trial, study)
 
-            # 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)
+def track_in_mlflow(optuna_mlflow: MLFlowIntegrator) -> Callable:
+    def decorator(func: Callable) -> Callable:
+        def wrapper(trial: optuna.trial.Trial) -> float:
+            study = trial.study
+            optuna_mlflow.set_experiment(study)
+
+            with mlflow.start_run(run_name=str(trial.number), nested=optuna_mlflow._nest_trials):
+                result = func(trial)
+                # This sets the metric for MLflow.
+                optuna_mlflow.log_metric(result)
+
+                # This sets the params for MLflow.
+                optuna_mlflow.log_params(trial)
+
+                # This sets the tags for MLflow.
+                optuna_mlflow.calculate_tags(trial, study)
+
+                return result
+
+        return wrapper
+
+    return decorator
diff --git a/tests/integration_tests/test_mlflow.py b/tests/integration_tests/test_mlflow.py
index 99e6acc92b9..2ae657cf4db 100644
--- a/tests/integration_tests/test_mlflow.py
+++ b/tests/integration_tests/test_mlflow.py
@@ -6,6 +6,8 @@
 
 import optuna
 from optuna.integration.mlflow import MLflowCallback
+from optuna.integration.mlflow import MLFlowIntegrator
+from optuna.integration.mlflow import track_in_mlflow
 
 
 def _objective_func(trial: optuna.trial.Trial) -> float:
@@ -210,3 +212,162 @@ 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_study_name_with_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
+
+    mlfli = MLFlowIntegrator(tracking_uri=tracking_file_name)
+
+    @track_in_mlflow(mlfli)
+    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)
+
+    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 "value" in first_run_dict["data"]["metrics"]
+    assert "x" in first_run_dict["data"]["params"]
+    assert "y" in first_run_dict["data"]["params"]
+    assert "z" in first_run_dict["data"]["params"]
+    assert first_run_dict["data"]["tags"]["direction"] == "MINIMIZE"
+    assert (
+        first_run_dict["data"]["tags"]["x_distribution"]
+        == "UniformDistribution(high=1.0, low=-1.0)"
+    )
+    assert (
+        first_run_dict["data"]["tags"]["y_distribution"]
+        == "LogUniformDistribution(high=30.0, low=20.0)"
+    )
+    assert (
+        first_run_dict["data"]["tags"]["z_distribution"]
+        == "CategoricalDistribution(choices=(-1.0, 1.0))"
+    )
+    assert first_run_dict["data"]["tags"]["my_user_attr"] == "my_user_attr_value"
+
+    assert metric_name in first_run_dict["data"]["metrics"]
+    assert first_run_dict["data"]["metrics"][metric_name] == metric
+
+
+def test_set_experiment(tmpdir: py.path.local) -> None:
+
+    tracking_file_name = "file:{}".format(tmpdir)
+    metric_name = "my_metric_name"
+    study_name = "my_study"
+
+    mlfli = MLFlowIntegrator(tracking_uri=tracking_file_name, metric_name=metric_name)
+    study = optuna.create_study(study_name=study_name)
+
+    mlfli.set_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
+
+    mlfli = MLFlowIntegrator(tracking_uri=tracking_file_name, metric_name=metric_name)
+    study = optuna.create_study(study_name=study_name)
+    mlfli.set_experiment(study)
+
+    with mlflow.start_run():
+        mlfli.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_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}
+
+    mlfli = MLFlowIntegrator(tracking_uri=tracking_file_name, metric_name=metric_name)
+    study = optuna.create_study(study_name=study_name)
+    mlfli.set_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,
+        )
+        mlfli.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)


Code review comments:
diff_hunk: @@ -13,8 +14,79 @@
     import mlflow
 
 
+class MLFlowIntegrator(object):
+    def __init__(
+        self,
+        tracking_uri: Optional[str] = None,
+        metric_name: str = "value",
+        nest_trials: bool = False,
+        tag_study_user_attrs: bool = False,
+    ) -> None:
+
+        _imports.check()
+
+        self._tracking_uri = tracking_uri
+        self._metric_name = metric_name
+        self._nest_trials = nest_trials
+        self._tag_study_user_attrs = tag_study_user_attrs
+
+    def set_experiment(self, study: optuna.study.Study) -> None:
path: optuna/integration/mlflow.py
text: @user1:
Could you rename this function? This is because `mlflow.set_tracking_uri` is also called in this function.

diff_hunk: @@ -85,73 +157,42 @@ def objective(trial):
             set, key value pairs in study.user_attrs will supersede existing tags.
     """
 
-    def __init__(
-        self,
-        tracking_uri: Optional[str] = None,
-        metric_name: str = "value",
-        nest_trials: bool = False,
-        tag_study_user_attrs: bool = False,
-    ) -> None:
-
-        _imports.check()
-
-        self._tracking_uri = tracking_uri
-        self._metric_name = metric_name
-        self._nest_trials = nest_trials
-        self._tag_study_user_attrs = tag_study_user_attrs
-
     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.set_experiment(study)
 
         with mlflow.start_run(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)
+            if trial.value:
path: optuna/integration/mlflow.py
text: @user1:
I feel this breaks a (minor) backward compatibility. Is it necessary this `if` statement? The tests seem fine without this line.

@author:
MLFlow does not like None as input for log_metric, if I remember right. I write a test for that (and if it is not necessary, remove the line)

@author:
I just checked, and the objective function actually should always return a float, hence we should be fine without this line. (Or actually a sequence of floats -> make it compatible with both.

I just checked with the original code, the None test was there already. Is it possible that the `trial_value` is `None` if the trial fails? 

@author:
I updated the code of `log_metric` such that it can handle `float`, `Sequence[float]` and `None`

diff_hunk: @@ -85,73 +157,42 @@ def objective(trial):
             set, key value pairs in study.user_attrs will supersede existing tags.
     """
 
-    def __init__(
-        self,
-        tracking_uri: Optional[str] = None,
-        metric_name: str = "value",
-        nest_trials: bool = False,
-        tag_study_user_attrs: bool = False,
-    ) -> None:
-
-        _imports.check()
-
-        self._tracking_uri = tracking_uri
-        self._metric_name = metric_name
-        self._nest_trials = nest_trials
-        self._tag_study_user_attrs = tag_study_user_attrs
-
     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.set_experiment(study)
 
         with mlflow.start_run(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)
+            if 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)
+            self.calculate_tags(trial, study)
 
-            # 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)
+def track_in_mlflow(optuna_mlflow: MLFlowIntegrator) -> Callable:
path: optuna/integration/mlflow.py
text: @user2:
Could you add `@user1("2.8.0")`?

diff_hunk: @@ -13,8 +14,79 @@
     import mlflow
 
 
+class MLFlowIntegrator(object):
+    def __init__(
+        self,
+        tracking_uri: Optional[str] = None,
+        metric_name: str = "value",
+        nest_trials: bool = False,
+        tag_study_user_attrs: bool = False,
+    ) -> None:
+
+        _imports.check()
+
+        self._tracking_uri = tracking_uri
+        self._metric_name = metric_name
+        self._nest_trials = nest_trials
+        self._tag_study_user_attrs = tag_study_user_attrs
+
+    def set_experiment(self, study: optuna.study.Study) -> 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)
+
+    def calculate_tags(self, trial: optuna.trial.BaseTrial, study: optuna.study.Study) -> None:
+
+        tags: Dict[str, str] = {}
+        tags["number"] = str(trial.number)
+        tags["datetime_start"] = str(trial.datetime_start)
+        if isinstance(trial, optuna.trial.FrozenTrial):
+            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
+        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)
+
+        # This sets the tags for MLflow.
+        mlflow.set_tags(tags)
+
+    def log_metric(self, value: float) -> None:
+        trial_value = value if value is not None else float("nan")
+        mlflow.log_metric(self._metric_name, trial_value)
+
+    def log_params(self, trial: optuna.trial.BaseTrial) -> None:
path: optuna/integration/mlflow.py
text: @user2:
Can we define this function as a static method: `add @user1` just before the definition?

diff_hunk: @@ -13,8 +14,79 @@
     import mlflow
 
 
+class MLFlowIntegrator(object):
path: optuna/integration/mlflow.py
text: @user2:
Could you add `@user1("2.8.0")`?

diff_hunk: @@ -13,8 +14,79 @@
     import mlflow
 
 
+class MLFlowIntegrator(object):
+    def __init__(
+        self,
+        tracking_uri: Optional[str] = None,
+        metric_name: str = "value",
+        nest_trials: bool = False,
+        tag_study_user_attrs: bool = False,
+    ) -> None:
+
+        _imports.check()
+
+        self._tracking_uri = tracking_uri
+        self._metric_name = metric_name
+        self._nest_trials = nest_trials
+        self._tag_study_user_attrs = tag_study_user_attrs
+
+    def set_experiment(self, study: optuna.study.Study) -> 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)
+
+    def calculate_tags(self, trial: optuna.trial.BaseTrial, study: optuna.study.Study) -> None:
+
+        tags: Dict[str, str] = {}
+        tags["number"] = str(trial.number)
+        tags["datetime_start"] = str(trial.datetime_start)
+        if isinstance(trial, optuna.trial.FrozenTrial):
path: optuna/integration/mlflow.py
text: @user1:
I might miss some points, but let me ask: why is this block for only `FrozenTrial`?

@author:
It is quite possible that I got something wrong. I have not been looking into the internals of Optuna before, but the property `datetime_complete` is not available on other types of trial after my observation.

@user1:
Thank you for your clarification. You are right! Related to the comment, can I ask why the trial's type is relaxed from `FrozenTrial` to `baseTrial` in the argument of `set_tags`?

@author:
`set_tags` is called from the Callback and the decorator.
The callback is called after the result of the optimization function has been stored by Optuna and the `Trial` has been transformed to a `FrozenTrial`. However in the context of the decorator `set_tag` is called before returning back to the Study and hence it is still a `Trial`. I thought that the most proper way to handle this difference for this two cases using this one instance check. As always I am open for other ideas.


Just as information MLFlow also tracks the duration of an trial/run. However when using the Callback functionalty this value is not correct since for MLFlow the run only starts in the callback

@author:
In the Callback we are working with a `FrozenTrial` (executed after the optimization function call has been handled by the study.
In the Decorator we did not return the to Study yet and hence still have a `Trial`.
That's why I had to relax to `BaseTrial` and introduce the instance check. 

If you have any better idea how to approach this, I am always open.

Just as a side note: MLFlow tracks to start and end of a run (=trial) as well. However when using the Callback functionality this information is not accurate since the run is only started after the trial has ended.

@user1:
Thank you for your clarification. So `@user2` cannot save`TrialState` and `datetime_complete` as in the tag attributes, and it cannot store a trial whose state is `PrunedTrial`.




@user1:
I think one way to resolve it is to resume the run; `track_in_mlflow` provides only `with` block to enable users to call mlflow's logging function and stores the `run_id`, then we resume the run by using `run_id` in the callback to perform mlflow's logging for `FrozenTrial`.

For example, https://github.com/nzw0301/optuna/commit/67943f1f2fd0f977043aaef5f96cf86314103e97

I believe there is room to improve the API, but what do you think of this?

@author:
I like it. It is less a break and more an extension of the current behaviour than my proposition
Two points I think should be discussed
- What happens when somebody does not add the Callback but sets the track_in_mlflow -> Optuna does not log anything (I can life with that)
- Multi-threaded behavior: What happens if the `nb_jobs` is set in the `optimize`-call. I know it is deprecated but still available until v4.

I am not that sure if if an attribute is the right place to store the run-id for me it looks dangerous for potential multi-threading and exception handling. My guts tell me it would be better to save it on the trial. Could `system_attrs` be used for that?

@user3:
Thanks for fruitful discussions @user1 and @author. I agree with @user1's idea. For the concerns by @author, I think
- If users set `track_in_mlflow` but the callback, Optuna should not log anything. It is sufficient to document the behavior enough.
- If the `n_jobs` argument is set, `MLFlowIntegrator` should raise an ValueError until v4. It is also sufficient to document the behavior.

> I am not that sure if if an attribute is the right place to store the run-id for me it looks dangerous for potential multi-threading and exception handling.

If we decide not to support the multi-threading with this MLFlow integration, we don't have to consider this problem, right? And, I'm not sure about the "exception handling". Could you elaborate a little more about it?

@author:
I have not thought all cases through, but since the proposed flow is
1. Set id in callback class in decorator
2. Read out id when in callback function
3. Reset id in callback class
Whenever something happens between 1 and 3 the id would not be reset. But since we would overwrite it in 1 again there would not be much harm to be done. 

But I still would prefer to store the information on the trial than in a "persistent" variable.

@author:
We @user1 approach by saving the run_id in a persistent variable in the class we have the following flow
1. Set run_id
2. Read run_id
Unless 1 is run again, the id is never reset. So if out of some error 1 is not set properly for the next trial, we overwrite the existing trial.

Other problematic pattern could be
```python
mlflc = MLFlowCallback(...)

@user4.track_in_mlflow()
def objective1(trial):
   mlflow.log_metric('something', 1)
   ...
   return x

def objective2(trial)
   ...
   return x*x

study1 = optuna.create_study(study_name="my_study")
study1.optimize(objective1, n_trials=10, callbacks=[mlflc])

study2 = optuna.create_study(study_name="my_study")
study2.optimize(objective2, n_trials=10, callbacks=[mlflc])
```
This might be far from best practice, but nothing prevents somebody to do this and the seconde study will log everything into the last run from the previous study.

I think it is possible to create other problematic scenarios with saving the current run into a variable of the class. This problems could be avoided if the run_id is not stored in MLFlowCallback but on the trial p.ex in `system_attrs` or `user_attrs`

@user1:
Thank you for your informative comments! As another scenario, we might want to resume optimization of the study like your example provided just before the comment of this comment with `optuna.study.load_study`. So personally, a re-usable callback sounds fine for some cases.

@author:
Ok this should be fine with any of the implementations. I was also thinking about the fact that in any cases it might be useful to have a reference to the MLFlow-Run-ID stored in Optuna to allow jumping from one to an other. Still not sure if it should be `system_attrs` or `user_attrs`

diff_hunk: @@ -13,8 +14,79 @@
     import mlflow
 
 
+class MLFlowIntegrator(object):
+    def __init__(
+        self,
+        tracking_uri: Optional[str] = None,
+        metric_name: str = "value",
+        nest_trials: bool = False,
+        tag_study_user_attrs: bool = False,
+    ) -> None:
+
+        _imports.check()
+
+        self._tracking_uri = tracking_uri
+        self._metric_name = metric_name
+        self._nest_trials = nest_trials
+        self._tag_study_user_attrs = tag_study_user_attrs
+
+    def set_experiment(self, study: optuna.study.Study) -> 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)
+
+    def calculate_tags(self, trial: optuna.trial.BaseTrial, study: optuna.study.Study) -> None:
path: optuna/integration/mlflow.py
text: @user1:
How about `set_tags` since this function does only tag-related things?