Repository: optuna/optuna

Issue:
MLflowCallback error when using n_jobs > 1
### Expected behavior

be able to use more than one job and at the same time log it in mlflow

### Environment

- Optuna version: 3.0.3
- Python version: 3.9.7
- OS: Linux-4.14.35-1902.8.4.el7uek.x86_64-x86_64-with-glibc2.31
- MLFlow version:  1.29.0


### Error messages, stack traces, or logs

```shell
Exception: Run with UUID 13a5f691281742d1a9d6b6f98c43203a is already active. To start a new run, first end the current run with mlflow.end_run(). To start a nested run, call start_run with nested=True
```


### Steps to reproduce

trying to use multiprocessing i get 

i know it is an mlflow issue, but there are some fixes to do when interacting with the mlflow api to support multiprocessing
so, some changes to MLflowCallback should be done

potential fixes
https://github.com/mlflow/mlflow/issues/3592
https://github.com/mlflow/mlflow/issues/1420

i tried even using:
`
    mlflow_kwargs={
        'nested': False
    }
`
but this raises duplication data and errors in mlflow

i also tried with
`
from joblib import parallel_backend
with parallel_backend('multiprocessing'):`
during study.optimize()

### Additional context (optional)

_No response_


Discussion:
Could you share the minimum reproducible code, please?
FYI: @toshihikoyanase 

```python
import optuna
import mlflow
from optuna.integration.mlflow import MLflowCallback

mlflc = MLflowCallback(
    metric_name="my metric score",
)


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=100, callbacks=[mlflc], n_jobs=4)
```
@nzw0301 Thank you for providing the reproducible code. The provided code causes the same error without `n_jobs=4`. I may miss something, but I guess the cause of the error may not be related to the parallelization.

```python
import optuna
import mlflow
from optuna.integration.mlflow import MLflowCallback

mlflc = MLflowCallback(
    metric_name="my metric score",
)


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=100, callbacks=[mlflc])
```

Without `mlflow.log_param` and `mlflow.log_metric`, the sequential optimization finished successfully, while the parallel optimization failed with the following error.

```console
mlflow.exceptions.MlflowException: Changing param values is not allowed. Param with key='x' was already logged with value='-1.4760283430008698' for run ID='b21577ccc7b94137a18aeb019526c023'. Attempted logging new value '-4.902930734265949'.
```

I guess the mlflow integration has multiple issues related to your reproducible code.
@toshihikoyanase Thank you for your comment. Sorry for confusing you by sharing the wrong code. The reproducible code should be as follows:


```python
import optuna
from optuna.integration.mlflow import MLflowCallback


def objective(trial):
    x = trial.suggest_float("x", -10, 10)
    return (x - 2) ** 2


mlflc = MLflowCallback(
#     tracking_uri=YOUR_TRACKING_URI,
    metric_name="my metric score",
)

study = optuna.create_study(study_name="my_study")
study.optimize(objective, n_trials=10, callbacks=[mlflc], n_jobs=4)
```

This code is almost the same as https://optuna.readthedocs.io/en/stable/reference/generated/optuna.integration.MLflowCallback.html.


Note that the first shared code came from MLFLowCallback's decorator version with the wrong modification...
@gen740 is working on this issue.
Optuna uses [mlflow fluent tracking API](https://mlflow.org/docs/latest/python_api/mlflow.html#mlflow.start_run), which is not threadsafe.

> The fluent tracking API is not currently threadsafe.

This bug is caused by race-condition of `mlflow.start_run()`.

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,4 +1,5 @@
 import functools
+import threading
 from typing import Any
 from typing import Callable
 from typing import Dict
@@ -139,27 +140,29 @@ def __init__(
         self._mlflow_kwargs = mlflow_kwargs or {}
         self._tag_study_user_attrs = tag_study_user_attrs
         self._tag_trial_user_attrs = tag_trial_user_attrs
+        self._lock = threading.Lock()
 
     def __call__(self, study: optuna.study.Study, trial: optuna.trial.FrozenTrial) -> None:
 
-        self._initialize_experiment(study)
+        with self._lock:
+            self._initialize_experiment(study)
 
-        with mlflow.start_run(
-            run_id=trial.system_attrs.get(RUN_ID_ATTRIBUTE_KEY),
-            experiment_id=self._mlflow_kwargs.get("experiment_id"),
-            run_name=self._mlflow_kwargs.get("run_name") or str(trial.number),
-            nested=self._mlflow_kwargs.get("nested") or False,
-            tags=self._mlflow_kwargs.get("tags"),
-        ):
+            with mlflow.start_run(
+                run_id=trial.system_attrs.get(RUN_ID_ATTRIBUTE_KEY),
+                experiment_id=self._mlflow_kwargs.get("experiment_id"),
+                run_name=self._mlflow_kwargs.get("run_name") or str(trial.number),
+                nested=self._mlflow_kwargs.get("nested") or False,
+                tags=self._mlflow_kwargs.get("tags"),
+            ):
 
-            # This sets the metrics for MLflow.
-            self._log_metrics(trial.values)
+                # This sets the metrics for MLflow.
+                self._log_metrics(trial.values)
 
-            # This sets the params for MLflow.
-            self._log_params(trial.params)
+                # This sets the params for MLflow.
+                self._log_params(trial.params)
 
-            # This sets the tags for MLflow.
-            self._set_tags(trial, study)
+                # This sets the tags for MLflow.
+                self._set_tags(trial, study)
 
     @experimental_func("2.9.0")
     def track_in_mlflow(self) -> Callable:
@@ -205,14 +208,15 @@ def objective(trial):
         def decorator(func: ObjectiveFuncType) -> ObjectiveFuncType:
             @functools.wraps(func)
             def wrapper(trial: optuna.trial.Trial) -> Union[float, Sequence[float]]:
-                study = trial.study
-                self._initialize_experiment(study)
-                nested = self._mlflow_kwargs.get("nested")
+                with self._lock:
+                    study = trial.study
+                    self._initialize_experiment(study)
+                    nested = self._mlflow_kwargs.get("nested")
 
-                with mlflow.start_run(run_name=str(trial.number), nested=nested) as run:
-                    trial.set_system_attr(RUN_ID_ATTRIBUTE_KEY, run.info.run_id)
+                    with mlflow.start_run(run_name=str(trial.number), nested=nested) as run:
+                        trial.set_system_attr(RUN_ID_ATTRIBUTE_KEY, run.info.run_id)
 
-                    return func(trial)
+                        return func(trial)
 
             return wrapper
 


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
@@ -268,6 +268,26 @@ def test_nest_trials(tmpdir: py.path.local) -> None:
     assert all(set(r.data.metrics.keys()) == {"value"} for r in child_runs)
 
 
+@pytest.mark.parametrize("n_jobs", [2, 4])
+def test_multiple_jobs(tmpdir: py.path.local, n_jobs: int) -> None:
+    tracking_uri = f"file:{tmpdir}"
+    study_name = "my_study"
+    # The race-condition usually happens after first trial for each job
+    n_trials = n_jobs * 2
+
+    mlflc = MLflowCallback(tracking_uri=tracking_uri)
+    study = optuna.create_study(study_name=study_name)
+    study.optimize(_objective_func, n_trials=n_trials, callbacks=[mlflc], n_jobs=n_jobs)
+
+    mlfl_client = MlflowClient(tracking_uri)
+    assert len(mlfl_client.list_experiments()) == 1
+
+    experiment = mlfl_client.list_experiments()[0]
+    runs = mlfl_client.list_run_infos(experiment.experiment_id)
+
+    assert len(runs) == n_trials
+
+
 def test_mlflow_callback_fails_when_nest_trials_is_false_and_active_run_exists(
     tmpdir: py.path.local,
 ) -> None:
@@ -373,11 +393,12 @@ def test_log_mlflow_tags(tmpdir: py.path.local) -> None:
     assert all([tags[key] == str(value) for key, value in expected_tags.items()])
 
 
-def test_track_in_mlflow_decorator(tmpdir: py.path.local) -> None:
+@pytest.mark.parametrize("n_jobs", [1, 2, 4])
+def test_track_in_mlflow_decorator(tmpdir: py.path.local, n_jobs: int) -> None:
 
     tracking_uri = f"file:{tmpdir}"
     study_name = "my_study"
-    n_trials = 3
+    n_trials = n_jobs * 2
 
     metric_name = "additional_metric"
     metric = 3.14
@@ -398,7 +419,7 @@ def _objective_func(trial: optuna.trial.Trial) -> float:
     tracked_objective = mlflc.track_in_mlflow()(_objective_func)
 
     study = optuna.create_study(study_name=study_name)
-    study.optimize(tracked_objective, n_trials=n_trials, callbacks=[mlflc])
+    study.optimize(tracked_objective, n_trials=n_trials, callbacks=[mlflc], n_jobs=n_jobs)
 
     mlfl_client = MlflowClient(tracking_uri)
     experiments = mlfl_client.list_experiments()