We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN 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_

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `optuna/integration/mlflow.py`
Content:
```
1 import functools
2 from typing import Any
3 from typing import Callable
4 from typing import Dict
5 from typing import List
6 from typing import Optional
7 from typing import Sequence
8 from typing import Union
9 
10 import optuna
11 from optuna._experimental import experimental_class
12 from optuna._experimental import experimental_func
13 from optuna._imports import try_import
14 from optuna.study.study import ObjectiveFuncType
15 
16 
17 with try_import() as _imports:
18     import mlflow
19 
20 RUN_ID_ATTRIBUTE_KEY = "mlflow_run_id"
21 
22 
23 @experimental_class("1.4.0")
24 class MLflowCallback:
25     """Callback to track Optuna trials with MLflow.
26 
27     This callback adds relevant information that is
28     tracked by Optuna to MLflow.
29 
30     Example:
31 
32         Add MLflow callback to Optuna optimization.
33 
34         .. testsetup::
35 
36             import pathlib
37             import tempfile
38 
39             tempdir = tempfile.mkdtemp()
40             YOUR_TRACKING_URI = pathlib.Path(tempdir).as_uri()
41 
42         .. testcode::
43 
44             import optuna
45             from optuna.integration.mlflow import MLflowCallback
46 
47 
48             def objective(trial):
49                 x = trial.suggest_float("x", -10, 10)
50                 return (x - 2) ** 2
51 
52 
53             mlflc = MLflowCallback(
54                 tracking_uri=YOUR_TRACKING_URI,
55                 metric_name="my metric score",
56             )
57 
58             study = optuna.create_study(study_name="my_study")
59             study.optimize(objective, n_trials=10, callbacks=[mlflc])
60 
61         .. testcleanup::
62 
63             import shutil
64 
65             shutil.rmtree(tempdir)
66 
67     Args:
68         tracking_uri:
69             The URI of the MLflow tracking server.
70 
71             Please refer to `mlflow.set_tracking_uri
72             <https://www.mlflow.org/docs/latest/python_api/mlflow.html#mlflow.set_tracking_uri>`_
73             for more details.
74         metric_name:
75             Name assigned to optimized metric. In case of multi-objective optimization,
76             list of names can be passed. Those names will be assigned
77             to metrics in the order returned by objective function.
78             If single name is provided, or this argument is left to default value,
79             it will be broadcasted to each objective with a number suffix in order
80             returned by objective function e.g. two objectives and default metric name
81             will be logged as ``value_0`` and ``value_1``. The number of metrics must be
82             the same as the number of values an objective function returns.
83         create_experiment:
84             When :obj:`True`, new MLflow experiment will be created for each optimization run,
85             named after the Optuna study. Setting this argument to :obj:`False` lets user run
86             optimization under existing experiment, set via `mlflow.set_experiment
87             <https://www.mlflow.org/docs/latest/python_api/mlflow.html#mlflow.get_tracking_uri>`_,
88             by passing ``experiment_id`` as one of ``mlflow_kwargs`` or under default MLflow
89             experiment, when no additional arguments are passed. Note that this argument
90             must be set to :obj:`False` when using Optuna with this callback within
91             Databricks Notebook.
92         mlflow_kwargs:
93             Set of arguments passed when initializing MLflow run.
94             Please refer to `MLflow API documentation
95             <https://www.mlflow.org/docs/latest/python_api/mlflow.html#mlflow.start_run>`_
96             for more details.
97 
98                 .. note::
99                     ``nest_trials`` argument added in v2.3.0 is a part of ``mlflow_kwargs``
100                     since v3.0.0. Anyone using ``nest_trials=True`` should migrate to
101                     ``mlflow_kwargs={"nested": True}`` to avoid raising :exc:`TypeError`.
102 
103         tag_study_user_attrs:
104             Flag indicating whether or not to add the study's user attrs
105             to the mlflow trial as tags. Please note that when this flag is
106             set, key value pairs in :attr:`~optuna.study.Study.user_attrs`
107             will supersede existing tags.
108         tag_trial_user_attrs:
109             Flag indicating whether or not to add the trial's user attrs
110             to the mlflow trial as tags. Please note that when both trial and
111             study user attributes are logged, the latter will supersede the former
112             in case of a collision.
113 
114 
115     """
116 
117     def __init__(
118         self,
119         tracking_uri: Optional[str] = None,
120         metric_name: Union[str, Sequence[str]] = "value",
121         create_experiment: bool = True,
122         mlflow_kwargs: Optional[Dict[str, Any]] = None,
123         tag_study_user_attrs: bool = False,
124         tag_trial_user_attrs: bool = True,
125     ) -> None:
126 
127         _imports.check()
128 
129         if not isinstance(metric_name, Sequence):
130             raise TypeError(
131                 "Expected metric_name to be string or sequence of strings, got {}.".format(
132                     type(metric_name)
133                 )
134             )
135 
136         self._tracking_uri = tracking_uri
137         self._metric_name = metric_name
138         self._create_experiment = create_experiment
139         self._mlflow_kwargs = mlflow_kwargs or {}
140         self._tag_study_user_attrs = tag_study_user_attrs
141         self._tag_trial_user_attrs = tag_trial_user_attrs
142 
143     def __call__(self, study: optuna.study.Study, trial: optuna.trial.FrozenTrial) -> None:
144 
145         self._initialize_experiment(study)
146 
147         with mlflow.start_run(
148             run_id=trial.system_attrs.get(RUN_ID_ATTRIBUTE_KEY),
149             experiment_id=self._mlflow_kwargs.get("experiment_id"),
150             run_name=self._mlflow_kwargs.get("run_name") or str(trial.number),
151             nested=self._mlflow_kwargs.get("nested") or False,
152             tags=self._mlflow_kwargs.get("tags"),
153         ):
154 
155             # This sets the metrics for MLflow.
156             self._log_metrics(trial.values)
157 
158             # This sets the params for MLflow.
159             self._log_params(trial.params)
160 
161             # This sets the tags for MLflow.
162             self._set_tags(trial, study)
163 
164     @experimental_func("2.9.0")
165     def track_in_mlflow(self) -> Callable:
166         """Decorator for using MLflow logging in the objective function.
167 
168         This decorator enables the extension of MLflow logging provided by the callback.
169 
170         All information logged in the decorated objective function will be added to the MLflow
171         run for the trial created by the callback.
172 
173         Example:
174 
175             Add additional logging to MLflow.
176 
177             .. testcode::
178 
179                 import optuna
180                 import mlflow
181                 from optuna.integration.mlflow import MLflowCallback
182 
183                 mlflc = MLflowCallback(
184                     tracking_uri=YOUR_TRACKING_URI,
185                     metric_name="my metric score",
186                 )
187 
188 
189                 @mlflc.track_in_mlflow()
190                 def objective(trial):
191                     x = trial.suggest_float("x", -10, 10)
192                     mlflow.log_param("power", 2)
193                     mlflow.log_metric("base of metric", x - 2)
194 
195                     return (x - 2) ** 2
196 
197 
198                 study = optuna.create_study(study_name="my_other_study")
199                 study.optimize(objective, n_trials=10, callbacks=[mlflc])
200 
201         Returns:
202             ObjectiveFuncType: Objective function with tracking to MLflow enabled.
203         """
204 
205         def decorator(func: ObjectiveFuncType) -> ObjectiveFuncType:
206             @functools.wraps(func)
207             def wrapper(trial: optuna.trial.Trial) -> Union[float, Sequence[float]]:
208                 study = trial.study
209                 self._initialize_experiment(study)
210                 nested = self._mlflow_kwargs.get("nested")
211 
212                 with mlflow.start_run(run_name=str(trial.number), nested=nested) as run:
213                     trial.set_system_attr(RUN_ID_ATTRIBUTE_KEY, run.info.run_id)
214 
215                     return func(trial)
216 
217             return wrapper
218 
219         return decorator
220 
221     def _initialize_experiment(self, study: optuna.study.Study) -> None:
222         """Initialize an MLflow experiment with the study name.
223 
224         If a tracking uri has been provided, MLflow will be initialized to use it.
225 
226         Args:
227             study: Study to be tracked in MLflow.
228         """
229 
230         # This sets the `tracking_uri` for MLflow.
231         if self._tracking_uri is not None:
232             mlflow.set_tracking_uri(self._tracking_uri)
233 
234         if self._create_experiment:
235             mlflow.set_experiment(study.study_name)
236 
237     def _set_tags(self, trial: optuna.trial.FrozenTrial, study: optuna.study.Study) -> None:
238         """Sets the Optuna tags for the current MLflow run.
239 
240         Args:
241             trial: Trial to be tracked.
242             study: Study to be tracked.
243         """
244 
245         tags: Dict[str, Union[str, List[str]]] = {}
246         tags["number"] = str(trial.number)
247         tags["datetime_start"] = str(trial.datetime_start)
248 
249         tags["datetime_complete"] = str(trial.datetime_complete)
250 
251         # Set trial state.
252         if trial.state.is_finished():
253             tags["state"] = trial.state.name
254 
255         # Set study directions.
256         directions = [d.name for d in study.directions]
257         tags["direction"] = directions if len(directions) != 1 else directions[0]
258 
259         distributions = {(k + "_distribution"): str(v) for (k, v) in trial.distributions.items()}
260         tags.update(distributions)
261 
262         if self._tag_trial_user_attrs:
263             tags.update(trial.user_attrs)
264 
265         if self._tag_study_user_attrs:
266             tags.update(study.user_attrs)
267 
268         # This is a temporary fix on Optuna side. It avoids an error with user
269         # attributes that are too long. It should be fixed on MLflow side later.
270         # When it is fixed on MLflow side this codeblock can be removed.
271         # see https://github.com/optuna/optuna/issues/1340
272         # see https://github.com/mlflow/mlflow/issues/2931
273         for key, value in tags.items():
274             value = str(value)  # make sure it is a string
275             max_val_length = mlflow.utils.validation.MAX_TAG_VAL_LENGTH
276             if len(value) > max_val_length:
277                 tags[key] = "{}...".format(value[: max_val_length - 3])
278 
279         mlflow.set_tags(tags)
280 
281     def _log_metrics(self, values: Optional[List[float]]) -> None:
282         """Log the trial results as metrics to MLflow.
283 
284         Args:
285             values: Results of a trial.
286         """
287         if values is None:
288             return
289 
290         if isinstance(self._metric_name, str):
291             if len(values) > 1:
292                 # Broadcast default name for multi-objective optimization.
293                 names = ["{}_{}".format(self._metric_name, i) for i in range(len(values))]
294 
295             else:
296                 names = [self._metric_name]
297 
298         else:
299             if len(self._metric_name) != len(values):
300                 raise ValueError(
301                     "Running multi-objective optimization "
302                     "with {} objective values, but {} names specified. "
303                     "Match objective values and names, or use default broadcasting.".format(
304                         len(values), len(self._metric_name)
305                     )
306                 )
307 
308             else:
309                 names = [*self._metric_name]
310 
311         metrics = {name: val for name, val in zip(names, values)}
312         mlflow.log_metrics(metrics)
313 
314     @staticmethod
315     def _log_params(params: Dict[str, Any]) -> None:
316         """Log the parameters of the trial to MLflow.
317 
318         Args:
319             params: Trial params.
320         """
321         mlflow.log_params(params)
322
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

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 @@
         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 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

{"golden_diff": "diff --git a/optuna/integration/mlflow.py b/optuna/integration/mlflow.py\n--- a/optuna/integration/mlflow.py\n+++ b/optuna/integration/mlflow.py\n@@ -1,4 +1,5 @@\n import functools\n+import threading\n from typing import Any\n from typing import Callable\n from typing import Dict\n@@ -139,27 +140,29 @@\n         self._mlflow_kwargs = mlflow_kwargs or {}\n         self._tag_study_user_attrs = tag_study_user_attrs\n         self._tag_trial_user_attrs = tag_trial_user_attrs\n+        self._lock = threading.Lock()\n \n     def __call__(self, study: optuna.study.Study, trial: optuna.trial.FrozenTrial) -> None:\n \n-        self._initialize_experiment(study)\n+        with self._lock:\n+            self._initialize_experiment(study)\n \n-        with mlflow.start_run(\n-            run_id=trial.system_attrs.get(RUN_ID_ATTRIBUTE_KEY),\n-            experiment_id=self._mlflow_kwargs.get(\"experiment_id\"),\n-            run_name=self._mlflow_kwargs.get(\"run_name\") or str(trial.number),\n-            nested=self._mlflow_kwargs.get(\"nested\") or False,\n-            tags=self._mlflow_kwargs.get(\"tags\"),\n-        ):\n+            with mlflow.start_run(\n+                run_id=trial.system_attrs.get(RUN_ID_ATTRIBUTE_KEY),\n+                experiment_id=self._mlflow_kwargs.get(\"experiment_id\"),\n+                run_name=self._mlflow_kwargs.get(\"run_name\") or str(trial.number),\n+                nested=self._mlflow_kwargs.get(\"nested\") or False,\n+                tags=self._mlflow_kwargs.get(\"tags\"),\n+            ):\n \n-            # This sets the metrics for MLflow.\n-            self._log_metrics(trial.values)\n+                # This sets the metrics for MLflow.\n+                self._log_metrics(trial.values)\n \n-            # This sets the params for MLflow.\n-            self._log_params(trial.params)\n+                # This sets the params for MLflow.\n+                self._log_params(trial.params)\n \n-            # This sets the tags for MLflow.\n-            self._set_tags(trial, study)\n+                # This sets the tags for MLflow.\n+                self._set_tags(trial, study)\n \n     @experimental_func(\"2.9.0\")\n     def track_in_mlflow(self) -> Callable:\n@@ -205,14 +208,15 @@\n         def decorator(func: ObjectiveFuncType) -> ObjectiveFuncType:\n             @functools.wraps(func)\n             def wrapper(trial: optuna.trial.Trial) -> Union[float, Sequence[float]]:\n-                study = trial.study\n-                self._initialize_experiment(study)\n-                nested = self._mlflow_kwargs.get(\"nested\")\n+                with self._lock:\n+                    study = trial.study\n+                    self._initialize_experiment(study)\n+                    nested = self._mlflow_kwargs.get(\"nested\")\n \n-                with mlflow.start_run(run_name=str(trial.number), nested=nested) as run:\n-                    trial.set_system_attr(RUN_ID_ATTRIBUTE_KEY, run.info.run_id)\n+                    with mlflow.start_run(run_name=str(trial.number), nested=nested) as run:\n+                        trial.set_system_attr(RUN_ID_ATTRIBUTE_KEY, run.info.run_id)\n \n-                    return func(trial)\n+                        return func(trial)\n \n             return wrapper\n", "issue": "MLflowCallback error when using n_jobs > 1\n### Expected behavior\r\n\r\nbe able to use more than one job and at the same time log it in mlflow\r\n\r\n### Environment\r\n\r\n- Optuna version: 3.0.3\r\n- Python version: 3.9.7\r\n- OS: Linux-4.14.35-1902.8.4.el7uek.x86_64-x86_64-with-glibc2.31\r\n- MLFlow version:  1.29.0\r\n\r\n\r\n### Error messages, stack traces, or logs\r\n\r\n```shell\r\nException: 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\r\n```\r\n\r\n\r\n### Steps to reproduce\r\n\r\ntrying to use multiprocessing i get \r\n\r\ni know it is an mlflow issue, but there are some fixes to do when interacting with the mlflow api to support multiprocessing\r\nso, some changes to MLflowCallback should be done\r\n\r\npotential fixes\r\nhttps://github.com/mlflow/mlflow/issues/3592\r\nhttps://github.com/mlflow/mlflow/issues/1420\r\n\r\ni tried even using:\r\n`\r\n    mlflow_kwargs={\r\n        'nested': False\r\n    }\r\n`\r\nbut this raises duplication data and errors in mlflow\r\n\r\ni also tried with\r\n`\r\nfrom joblib import parallel_backend\r\nwith parallel_backend('multiprocessing'):`\r\nduring study.optimize()\r\n\r\n### Additional context (optional)\r\n\r\n_No response_\n", "before_files": [{"content": "import functools\nfrom typing import Any\nfrom typing import Callable\nfrom typing import Dict\nfrom typing import List\nfrom typing import Optional\nfrom typing import Sequence\nfrom typing import Union\n\nimport optuna\nfrom optuna._experimental import experimental_class\nfrom optuna._experimental import experimental_func\nfrom optuna._imports import try_import\nfrom optuna.study.study import ObjectiveFuncType\n\n\nwith try_import() as _imports:\n    import mlflow\n\nRUN_ID_ATTRIBUTE_KEY = \"mlflow_run_id\"\n\n\n@experimental_class(\"1.4.0\")\nclass MLflowCallback:\n    \"\"\"Callback to track Optuna trials with MLflow.\n\n    This callback adds relevant information that is\n    tracked by Optuna to MLflow.\n\n    Example:\n\n        Add MLflow callback to Optuna optimization.\n\n        .. testsetup::\n\n            import pathlib\n            import tempfile\n\n            tempdir = tempfile.mkdtemp()\n            YOUR_TRACKING_URI = pathlib.Path(tempdir).as_uri()\n\n        .. testcode::\n\n            import optuna\n            from optuna.integration.mlflow import MLflowCallback\n\n\n            def objective(trial):\n                x = trial.suggest_float(\"x\", -10, 10)\n                return (x - 2) ** 2\n\n\n            mlflc = MLflowCallback(\n                tracking_uri=YOUR_TRACKING_URI,\n                metric_name=\"my metric score\",\n            )\n\n            study = optuna.create_study(study_name=\"my_study\")\n            study.optimize(objective, n_trials=10, callbacks=[mlflc])\n\n        .. testcleanup::\n\n            import shutil\n\n            shutil.rmtree(tempdir)\n\n    Args:\n        tracking_uri:\n            The URI of the MLflow tracking server.\n\n            Please refer to `mlflow.set_tracking_uri\n            <https://www.mlflow.org/docs/latest/python_api/mlflow.html#mlflow.set_tracking_uri>`_\n            for more details.\n        metric_name:\n            Name assigned to optimized metric. In case of multi-objective optimization,\n            list of names can be passed. Those names will be assigned\n            to metrics in the order returned by objective function.\n            If single name is provided, or this argument is left to default value,\n            it will be broadcasted to each objective with a number suffix in order\n            returned by objective function e.g. two objectives and default metric name\n            will be logged as ``value_0`` and ``value_1``. The number of metrics must be\n            the same as the number of values an objective function returns.\n        create_experiment:\n            When :obj:`True`, new MLflow experiment will be created for each optimization run,\n            named after the Optuna study. Setting this argument to :obj:`False` lets user run\n            optimization under existing experiment, set via `mlflow.set_experiment\n            <https://www.mlflow.org/docs/latest/python_api/mlflow.html#mlflow.get_tracking_uri>`_,\n            by passing ``experiment_id`` as one of ``mlflow_kwargs`` or under default MLflow\n            experiment, when no additional arguments are passed. Note that this argument\n            must be set to :obj:`False` when using Optuna with this callback within\n            Databricks Notebook.\n        mlflow_kwargs:\n            Set of arguments passed when initializing MLflow run.\n            Please refer to `MLflow API documentation\n            <https://www.mlflow.org/docs/latest/python_api/mlflow.html#mlflow.start_run>`_\n            for more details.\n\n                .. note::\n                    ``nest_trials`` argument added in v2.3.0 is a part of ``mlflow_kwargs``\n                    since v3.0.0. Anyone using ``nest_trials=True`` should migrate to\n                    ``mlflow_kwargs={\"nested\": True}`` to avoid raising :exc:`TypeError`.\n\n        tag_study_user_attrs:\n            Flag indicating whether or not to add the study's user attrs\n            to the mlflow trial as tags. Please note that when this flag is\n            set, key value pairs in :attr:`~optuna.study.Study.user_attrs`\n            will supersede existing tags.\n        tag_trial_user_attrs:\n            Flag indicating whether or not to add the trial's user attrs\n            to the mlflow trial as tags. Please note that when both trial and\n            study user attributes are logged, the latter will supersede the former\n            in case of a collision.\n\n\n    \"\"\"\n\n    def __init__(\n        self,\n        tracking_uri: Optional[str] = None,\n        metric_name: Union[str, Sequence[str]] = \"value\",\n        create_experiment: bool = True,\n        mlflow_kwargs: Optional[Dict[str, Any]] = None,\n        tag_study_user_attrs: bool = False,\n        tag_trial_user_attrs: bool = True,\n    ) -> None:\n\n        _imports.check()\n\n        if not isinstance(metric_name, Sequence):\n            raise TypeError(\n                \"Expected metric_name to be string or sequence of strings, got {}.\".format(\n                    type(metric_name)\n                )\n            )\n\n        self._tracking_uri = tracking_uri\n        self._metric_name = metric_name\n        self._create_experiment = create_experiment\n        self._mlflow_kwargs = mlflow_kwargs or {}\n        self._tag_study_user_attrs = tag_study_user_attrs\n        self._tag_trial_user_attrs = tag_trial_user_attrs\n\n    def __call__(self, study: optuna.study.Study, trial: optuna.trial.FrozenTrial) -> None:\n\n        self._initialize_experiment(study)\n\n        with mlflow.start_run(\n            run_id=trial.system_attrs.get(RUN_ID_ATTRIBUTE_KEY),\n            experiment_id=self._mlflow_kwargs.get(\"experiment_id\"),\n            run_name=self._mlflow_kwargs.get(\"run_name\") or str(trial.number),\n            nested=self._mlflow_kwargs.get(\"nested\") or False,\n            tags=self._mlflow_kwargs.get(\"tags\"),\n        ):\n\n            # This sets the metrics for MLflow.\n            self._log_metrics(trial.values)\n\n            # This sets the params for MLflow.\n            self._log_params(trial.params)\n\n            # This sets the tags for MLflow.\n            self._set_tags(trial, study)\n\n    @experimental_func(\"2.9.0\")\n    def track_in_mlflow(self) -> Callable:\n        \"\"\"Decorator for using MLflow logging in the objective function.\n\n        This decorator enables the extension of MLflow logging provided by the callback.\n\n        All information logged in the decorated objective function will be added to the MLflow\n        run for the trial created by the callback.\n\n        Example:\n\n            Add additional logging to MLflow.\n\n            .. testcode::\n\n                import optuna\n                import mlflow\n                from optuna.integration.mlflow import MLflowCallback\n\n                mlflc = MLflowCallback(\n                    tracking_uri=YOUR_TRACKING_URI,\n                    metric_name=\"my metric score\",\n                )\n\n\n                @mlflc.track_in_mlflow()\n                def objective(trial):\n                    x = trial.suggest_float(\"x\", -10, 10)\n                    mlflow.log_param(\"power\", 2)\n                    mlflow.log_metric(\"base of metric\", x - 2)\n\n                    return (x - 2) ** 2\n\n\n                study = optuna.create_study(study_name=\"my_other_study\")\n                study.optimize(objective, n_trials=10, callbacks=[mlflc])\n\n        Returns:\n            ObjectiveFuncType: Objective function with tracking to MLflow enabled.\n        \"\"\"\n\n        def decorator(func: ObjectiveFuncType) -> ObjectiveFuncType:\n            @functools.wraps(func)\n            def wrapper(trial: optuna.trial.Trial) -> Union[float, Sequence[float]]:\n                study = trial.study\n                self._initialize_experiment(study)\n                nested = self._mlflow_kwargs.get(\"nested\")\n\n                with mlflow.start_run(run_name=str(trial.number), nested=nested) as run:\n                    trial.set_system_attr(RUN_ID_ATTRIBUTE_KEY, run.info.run_id)\n\n                    return func(trial)\n\n            return wrapper\n\n        return decorator\n\n    def _initialize_experiment(self, study: optuna.study.Study) -> None:\n        \"\"\"Initialize an MLflow experiment with the study name.\n\n        If a tracking uri has been provided, MLflow will be initialized to use it.\n\n        Args:\n            study: Study to be tracked in MLflow.\n        \"\"\"\n\n        # This sets the `tracking_uri` for MLflow.\n        if self._tracking_uri is not None:\n            mlflow.set_tracking_uri(self._tracking_uri)\n\n        if self._create_experiment:\n            mlflow.set_experiment(study.study_name)\n\n    def _set_tags(self, trial: optuna.trial.FrozenTrial, study: optuna.study.Study) -> None:\n        \"\"\"Sets the Optuna tags for the current MLflow run.\n\n        Args:\n            trial: Trial to be tracked.\n            study: Study to be tracked.\n        \"\"\"\n\n        tags: Dict[str, Union[str, List[str]]] = {}\n        tags[\"number\"] = str(trial.number)\n        tags[\"datetime_start\"] = str(trial.datetime_start)\n\n        tags[\"datetime_complete\"] = str(trial.datetime_complete)\n\n        # Set trial state.\n        if trial.state.is_finished():\n            tags[\"state\"] = trial.state.name\n\n        # Set study directions.\n        directions = [d.name for d in study.directions]\n        tags[\"direction\"] = directions if len(directions) != 1 else directions[0]\n\n        distributions = {(k + \"_distribution\"): str(v) for (k, v) in trial.distributions.items()}\n        tags.update(distributions)\n\n        if self._tag_trial_user_attrs:\n            tags.update(trial.user_attrs)\n\n        if self._tag_study_user_attrs:\n            tags.update(study.user_attrs)\n\n        # This is a temporary fix on Optuna side. It avoids an error with user\n        # attributes that are too long. It should be fixed on MLflow side later.\n        # When it is fixed on MLflow side this codeblock can be removed.\n        # see https://github.com/optuna/optuna/issues/1340\n        # see https://github.com/mlflow/mlflow/issues/2931\n        for key, value in tags.items():\n            value = str(value)  # make sure it is a string\n            max_val_length = mlflow.utils.validation.MAX_TAG_VAL_LENGTH\n            if len(value) > max_val_length:\n                tags[key] = \"{}...\".format(value[: max_val_length - 3])\n\n        mlflow.set_tags(tags)\n\n    def _log_metrics(self, values: Optional[List[float]]) -> None:\n        \"\"\"Log the trial results as metrics to MLflow.\n\n        Args:\n            values: Results of a trial.\n        \"\"\"\n        if values is None:\n            return\n\n        if isinstance(self._metric_name, str):\n            if len(values) > 1:\n                # Broadcast default name for multi-objective optimization.\n                names = [\"{}_{}\".format(self._metric_name, i) for i in range(len(values))]\n\n            else:\n                names = [self._metric_name]\n\n        else:\n            if len(self._metric_name) != len(values):\n                raise ValueError(\n                    \"Running multi-objective optimization \"\n                    \"with {} objective values, but {} names specified. \"\n                    \"Match objective values and names, or use default broadcasting.\".format(\n                        len(values), len(self._metric_name)\n                    )\n                )\n\n            else:\n                names = [*self._metric_name]\n\n        metrics = {name: val for name, val in zip(names, values)}\n        mlflow.log_metrics(metrics)\n\n    @staticmethod\n    def _log_params(params: Dict[str, Any]) -> None:\n        \"\"\"Log the parameters of the trial to MLflow.\n\n        Args:\n            params: Trial params.\n        \"\"\"\n        mlflow.log_params(params)\n", "path": "optuna/integration/mlflow.py"}], "after_files": [{"content": "import functools\nimport threading\nfrom typing import Any\nfrom typing import Callable\nfrom typing import Dict\nfrom typing import List\nfrom typing import Optional\nfrom typing import Sequence\nfrom typing import Union\n\nimport optuna\nfrom optuna._experimental import experimental_class\nfrom optuna._experimental import experimental_func\nfrom optuna._imports import try_import\nfrom optuna.study.study import ObjectiveFuncType\n\n\nwith try_import() as _imports:\n    import mlflow\n\nRUN_ID_ATTRIBUTE_KEY = \"mlflow_run_id\"\n\n\n@experimental_class(\"1.4.0\")\nclass MLflowCallback:\n    \"\"\"Callback to track Optuna trials with MLflow.\n\n    This callback adds relevant information that is\n    tracked by Optuna to MLflow.\n\n    Example:\n\n        Add MLflow callback to Optuna optimization.\n\n        .. testsetup::\n\n            import pathlib\n            import tempfile\n\n            tempdir = tempfile.mkdtemp()\n            YOUR_TRACKING_URI = pathlib.Path(tempdir).as_uri()\n\n        .. testcode::\n\n            import optuna\n            from optuna.integration.mlflow import MLflowCallback\n\n\n            def objective(trial):\n                x = trial.suggest_float(\"x\", -10, 10)\n                return (x - 2) ** 2\n\n\n            mlflc = MLflowCallback(\n                tracking_uri=YOUR_TRACKING_URI,\n                metric_name=\"my metric score\",\n            )\n\n            study = optuna.create_study(study_name=\"my_study\")\n            study.optimize(objective, n_trials=10, callbacks=[mlflc])\n\n        .. testcleanup::\n\n            import shutil\n\n            shutil.rmtree(tempdir)\n\n    Args:\n        tracking_uri:\n            The URI of the MLflow tracking server.\n\n            Please refer to `mlflow.set_tracking_uri\n            <https://www.mlflow.org/docs/latest/python_api/mlflow.html#mlflow.set_tracking_uri>`_\n            for more details.\n        metric_name:\n            Name assigned to optimized metric. In case of multi-objective optimization,\n            list of names can be passed. Those names will be assigned\n            to metrics in the order returned by objective function.\n            If single name is provided, or this argument is left to default value,\n            it will be broadcasted to each objective with a number suffix in order\n            returned by objective function e.g. two objectives and default metric name\n            will be logged as ``value_0`` and ``value_1``. The number of metrics must be\n            the same as the number of values an objective function returns.\n        create_experiment:\n            When :obj:`True`, new MLflow experiment will be created for each optimization run,\n            named after the Optuna study. Setting this argument to :obj:`False` lets user run\n            optimization under existing experiment, set via `mlflow.set_experiment\n            <https://www.mlflow.org/docs/latest/python_api/mlflow.html#mlflow.get_tracking_uri>`_,\n            by passing ``experiment_id`` as one of ``mlflow_kwargs`` or under default MLflow\n            experiment, when no additional arguments are passed. Note that this argument\n            must be set to :obj:`False` when using Optuna with this callback within\n            Databricks Notebook.\n        mlflow_kwargs:\n            Set of arguments passed when initializing MLflow run.\n            Please refer to `MLflow API documentation\n            <https://www.mlflow.org/docs/latest/python_api/mlflow.html#mlflow.start_run>`_\n            for more details.\n\n                .. note::\n                    ``nest_trials`` argument added in v2.3.0 is a part of ``mlflow_kwargs``\n                    since v3.0.0. Anyone using ``nest_trials=True`` should migrate to\n                    ``mlflow_kwargs={\"nested\": True}`` to avoid raising :exc:`TypeError`.\n\n        tag_study_user_attrs:\n            Flag indicating whether or not to add the study's user attrs\n            to the mlflow trial as tags. Please note that when this flag is\n            set, key value pairs in :attr:`~optuna.study.Study.user_attrs`\n            will supersede existing tags.\n        tag_trial_user_attrs:\n            Flag indicating whether or not to add the trial's user attrs\n            to the mlflow trial as tags. Please note that when both trial and\n            study user attributes are logged, the latter will supersede the former\n            in case of a collision.\n\n\n    \"\"\"\n\n    def __init__(\n        self,\n        tracking_uri: Optional[str] = None,\n        metric_name: Union[str, Sequence[str]] = \"value\",\n        create_experiment: bool = True,\n        mlflow_kwargs: Optional[Dict[str, Any]] = None,\n        tag_study_user_attrs: bool = False,\n        tag_trial_user_attrs: bool = True,\n    ) -> None:\n\n        _imports.check()\n\n        if not isinstance(metric_name, Sequence):\n            raise TypeError(\n                \"Expected metric_name to be string or sequence of strings, got {}.\".format(\n                    type(metric_name)\n                )\n            )\n\n        self._tracking_uri = tracking_uri\n        self._metric_name = metric_name\n        self._create_experiment = create_experiment\n        self._mlflow_kwargs = mlflow_kwargs or {}\n        self._tag_study_user_attrs = tag_study_user_attrs\n        self._tag_trial_user_attrs = tag_trial_user_attrs\n        self._lock = threading.Lock()\n\n    def __call__(self, study: optuna.study.Study, trial: optuna.trial.FrozenTrial) -> None:\n\n        with self._lock:\n            self._initialize_experiment(study)\n\n            with mlflow.start_run(\n                run_id=trial.system_attrs.get(RUN_ID_ATTRIBUTE_KEY),\n                experiment_id=self._mlflow_kwargs.get(\"experiment_id\"),\n                run_name=self._mlflow_kwargs.get(\"run_name\") or str(trial.number),\n                nested=self._mlflow_kwargs.get(\"nested\") or False,\n                tags=self._mlflow_kwargs.get(\"tags\"),\n            ):\n\n                # This sets the metrics for MLflow.\n                self._log_metrics(trial.values)\n\n                # This sets the params for MLflow.\n                self._log_params(trial.params)\n\n                # This sets the tags for MLflow.\n                self._set_tags(trial, study)\n\n    @experimental_func(\"2.9.0\")\n    def track_in_mlflow(self) -> Callable:\n        \"\"\"Decorator for using MLflow logging in the objective function.\n\n        This decorator enables the extension of MLflow logging provided by the callback.\n\n        All information logged in the decorated objective function will be added to the MLflow\n        run for the trial created by the callback.\n\n        Example:\n\n            Add additional logging to MLflow.\n\n            .. testcode::\n\n                import optuna\n                import mlflow\n                from optuna.integration.mlflow import MLflowCallback\n\n                mlflc = MLflowCallback(\n                    tracking_uri=YOUR_TRACKING_URI,\n                    metric_name=\"my metric score\",\n                )\n\n\n                @mlflc.track_in_mlflow()\n                def objective(trial):\n                    x = trial.suggest_float(\"x\", -10, 10)\n                    mlflow.log_param(\"power\", 2)\n                    mlflow.log_metric(\"base of metric\", x - 2)\n\n                    return (x - 2) ** 2\n\n\n                study = optuna.create_study(study_name=\"my_other_study\")\n                study.optimize(objective, n_trials=10, callbacks=[mlflc])\n\n        Returns:\n            ObjectiveFuncType: Objective function with tracking to MLflow enabled.\n        \"\"\"\n\n        def decorator(func: ObjectiveFuncType) -> ObjectiveFuncType:\n            @functools.wraps(func)\n            def wrapper(trial: optuna.trial.Trial) -> Union[float, Sequence[float]]:\n                with self._lock:\n                    study = trial.study\n                    self._initialize_experiment(study)\n                    nested = self._mlflow_kwargs.get(\"nested\")\n\n                    with mlflow.start_run(run_name=str(trial.number), nested=nested) as run:\n                        trial.set_system_attr(RUN_ID_ATTRIBUTE_KEY, run.info.run_id)\n\n                        return func(trial)\n\n            return wrapper\n\n        return decorator\n\n    def _initialize_experiment(self, study: optuna.study.Study) -> None:\n        \"\"\"Initialize an MLflow experiment with the study name.\n\n        If a tracking uri has been provided, MLflow will be initialized to use it.\n\n        Args:\n            study: Study to be tracked in MLflow.\n        \"\"\"\n\n        # This sets the `tracking_uri` for MLflow.\n        if self._tracking_uri is not None:\n            mlflow.set_tracking_uri(self._tracking_uri)\n\n        if self._create_experiment:\n            mlflow.set_experiment(study.study_name)\n\n    def _set_tags(self, trial: optuna.trial.FrozenTrial, study: optuna.study.Study) -> None:\n        \"\"\"Sets the Optuna tags for the current MLflow run.\n\n        Args:\n            trial: Trial to be tracked.\n            study: Study to be tracked.\n        \"\"\"\n\n        tags: Dict[str, Union[str, List[str]]] = {}\n        tags[\"number\"] = str(trial.number)\n        tags[\"datetime_start\"] = str(trial.datetime_start)\n\n        tags[\"datetime_complete\"] = str(trial.datetime_complete)\n\n        # Set trial state.\n        if trial.state.is_finished():\n            tags[\"state\"] = trial.state.name\n\n        # Set study directions.\n        directions = [d.name for d in study.directions]\n        tags[\"direction\"] = directions if len(directions) != 1 else directions[0]\n\n        distributions = {(k + \"_distribution\"): str(v) for (k, v) in trial.distributions.items()}\n        tags.update(distributions)\n\n        if self._tag_trial_user_attrs:\n            tags.update(trial.user_attrs)\n\n        if self._tag_study_user_attrs:\n            tags.update(study.user_attrs)\n\n        # This is a temporary fix on Optuna side. It avoids an error with user\n        # attributes that are too long. It should be fixed on MLflow side later.\n        # When it is fixed on MLflow side this codeblock can be removed.\n        # see https://github.com/optuna/optuna/issues/1340\n        # see https://github.com/mlflow/mlflow/issues/2931\n        for key, value in tags.items():\n            value = str(value)  # make sure it is a string\n            max_val_length = mlflow.utils.validation.MAX_TAG_VAL_LENGTH\n            if len(value) > max_val_length:\n                tags[key] = \"{}...\".format(value[: max_val_length - 3])\n\n        mlflow.set_tags(tags)\n\n    def _log_metrics(self, values: Optional[List[float]]) -> None:\n        \"\"\"Log the trial results as metrics to MLflow.\n\n        Args:\n            values: Results of a trial.\n        \"\"\"\n        if values is None:\n            return\n\n        if isinstance(self._metric_name, str):\n            if len(values) > 1:\n                # Broadcast default name for multi-objective optimization.\n                names = [\"{}_{}\".format(self._metric_name, i) for i in range(len(values))]\n\n            else:\n                names = [self._metric_name]\n\n        else:\n            if len(self._metric_name) != len(values):\n                raise ValueError(\n                    \"Running multi-objective optimization \"\n                    \"with {} objective values, but {} names specified. \"\n                    \"Match objective values and names, or use default broadcasting.\".format(\n                        len(values), len(self._metric_name)\n                    )\n                )\n\n            else:\n                names = [*self._metric_name]\n\n        metrics = {name: val for name, val in zip(names, values)}\n        mlflow.log_metrics(metrics)\n\n    @staticmethod\n    def _log_params(params: Dict[str, Any]) -> None:\n        \"\"\"Log the parameters of the trial to MLflow.\n\n        Args:\n            params: Trial params.\n        \"\"\"\n        mlflow.log_params(params)\n", "path": "optuna/integration/mlflow.py"}]}