Repository: optuna/optuna

Issue:
Allow users to call `study.optimize()` in multiple threads.
### Motivation

For users who are familiar with `concurrent.futures`, I'd like to support the following usage that calls `study.optimize()` in multiple threads, although we can concurrently evaluate objective function in multiple threads via the `n_jobs` option of `study.optimize()`.

```python
import optuna
from concurrent.futures import ThreadPoolExecutor


def objective(trial):
    x = trial.suggest_float("x", -100, 100)
    y = trial.suggest_float("y", -100, 100)
    return x**2 + y


if __name__ == "__main__":
    study = optuna.create_study()
    with ThreadPoolExecutor(10) as pool:
        for i in range(10):
            pool.submit(study.optimize, objective, n_trials=10)

    print(f"Number of trials: {len(study.get_trials())}")
    print(f"Best params: {study.best_params}")
```

Also I'd like to support the following usage of `dask.distributed` (refs https://github.com/optuna/optuna/pull/2023#issuecomment-1268792200).


```python
import optuna

from dask.distributed import Client, wait


def objective(trial):
    x = trial.suggest_float("x", -100, 100)
    y = trial.suggest_float("y", -100, 100)
    return x**2 + y


if __name__ == "__main__":
    with Client("...") as client:
        print(f"Dask dashboard is available at {client.dashboard_link}")
        backend_storage = optuna.storages.InMemoryStorage()
        dask_storage = optuna.integration.DaskStorage(storage=backend_storage)
        study = optuna.create_study(direction="maximize", storage=dask_storage)

        futures = []
        for i in range(10):
            futures.append(client.submit(study.optimize, objective, n_trials=100, pure=False))
        wait(futures)

        print(f"Number of trials: {len(study.get_trials())}")
        print(f"Best params: {study.best_params}")
```

### Description

In the current implementation of Optuna, we can call `study.optimize()` in multiple processes like:

```python
import optuna
from concurrent.futures import ProcessPoolExecutor


def objective(trial):
    x = trial.suggest_float("x", -100, 100)
    y = trial.suggest_float("y", -100, 100)
    return x**2 + y


if __name__ == "__main__":
    storage = optuna.storages.RDBStorage(
        "sqlite:///example-db.sqlite3",
        engine_kwargs={"connect_args": {"timeout": 15.0}}
    )
    study = optuna.create_study(storage=storage)
    with ProcessPoolExecutor() as pool:
        for i in range(10):
            pool.submit(study.optimize, objective, n_trials=10)

    print(len(study.get_trials()))
```

However, it will be actually broken if we replace `ProcessPoolExecutor` with `ThreadPoolExecutor` due to following checks.

https://github.com/optuna/optuna/blob/6fc5392c458101d5bfda12d3f97dd88e11d4fd5b/optuna/study/_optimize.py#L53-L54
https://github.com/optuna/optuna/blob/6fc5392c458101d5bfda12d3f97dd88e11d4fd5b/optuna/study/_optimize.py#L121

As discussed at https://github.com/optuna/optuna/pull/2023#issuecomment-1268792200, the above checks are added at #491. In my recognition, it is originally intended to prevent the illegal calls of `study.optimize()`.
For example, Optuna's Trial object holds the Study object like [this](https://github.com/optuna/optuna/blob/master/optuna/trial/_trial.py#L26-L48), so components such as Sampler, Pruner, and the objective function can access the study through trials. It allows illegal calls from within the objective function like:

```python
def objective(trial):
    trial.study.optimize(objective, n_trials=10)
```

At least, I don't think it was intended to make the above use cases impossible. The above use cases will be fixed by simply removing the above three lines, but it may not be that easy as some features such as `study.stop()` and progress_bar may break.

### Alternatives (optional)

_No response_

### Additional context (optional)

_No response_


Discussion:

Resolving patch:
diff --git a/optuna/study/_optimize.py b/optuna/study/_optimize.py
--- a/optuna/study/_optimize.py
+++ b/optuna/study/_optimize.py
@@ -50,7 +50,7 @@ def _optimize(
             "The catch argument is of type '{}' but must be a tuple.".format(type(catch).__name__)
         )
 
-    if not study._optimize_lock.acquire(False):
+    if study._thread_local.in_optimize_loop:
         raise RuntimeError("Nested invocation of `Study.optimize` method isn't allowed.")
 
     if show_progress_bar and n_trials is None and timeout is not None and n_jobs != 1:
@@ -118,7 +118,7 @@ def _optimize(
                         )
                     )
     finally:
-        study._optimize_lock.release()
+        study._thread_local.in_optimize_loop = False
         progress_bar.close()
 
 
@@ -134,6 +134,9 @@ def _optimize_sequential(
     time_start: Optional[datetime.datetime],
     progress_bar: Optional[pbar_module._ProgressBar],
 ) -> None:
+    # Here we set `in_optimize_loop = True`, not at the beginning of the `_optimize()` function.
+    # Because it is a thread-local object and `n_jobs` option spawns new threads.
+    study._thread_local.in_optimize_loop = True
     if reseed_sampler_rng:
         study.sampler.reseed_rng()
 
diff --git a/optuna/study/study.py b/optuna/study/study.py
--- a/optuna/study/study.py
+++ b/optuna/study/study.py
@@ -53,6 +53,10 @@
 _logger = logging.get_logger(__name__)
 
 
+class _ThreadLocalStudyAttribute(threading.local):
+    in_optimize_loop: bool = False
+
+
 class Study:
     """A study corresponds to an optimization task, i.e., a set of trials.
 
@@ -82,19 +86,19 @@ def __init__(
         self.sampler = sampler or samplers.TPESampler()
         self.pruner = pruner or pruners.MedianPruner()
 
-        self._optimize_lock = threading.Lock()
+        self._thread_local = _ThreadLocalStudyAttribute()
         self._stop_flag = False
 
     def __getstate__(self) -> Dict[Any, Any]:
 
         state = self.__dict__.copy()
-        del state["_optimize_lock"]
+        del state["_thread_local"]
         return state
 
     def __setstate__(self, state: Dict[Any, Any]) -> None:
 
         self.__dict__.update(state)
-        self._optimize_lock = threading.Lock()
+        self._thread_local = _ThreadLocalStudyAttribute()
 
     @property
     def best_params(self) -> Dict[str, Any]:
@@ -494,9 +498,8 @@ def ask(
             A :class:`~optuna.trial.Trial`.
         """
 
-        if not self._optimize_lock.locked():
-            if is_heartbeat_enabled(self._storage):
-                warnings.warn("Heartbeat of storage is supposed to be used with Study.optimize.")
+        if not self._thread_local.in_optimize_loop and is_heartbeat_enabled(self._storage):
+            warnings.warn("Heartbeat of storage is supposed to be used with Study.optimize.")
 
         fixed_distributions = fixed_distributions or {}
         fixed_distributions = {
@@ -760,8 +763,7 @@ def objective(trial):
 
         """
 
-        if self._optimize_lock.acquire(False):
-            self._optimize_lock.release()
+        if not self._thread_local.in_optimize_loop:
             raise RuntimeError(
                 "`Study.stop` is supposed to be invoked inside an objective function or a "
                 "callback."


Test patch:
diff --git a/tests/study_tests/test_study.py b/tests/study_tests/test_study.py
--- a/tests/study_tests/test_study.py
+++ b/tests/study_tests/test_study.py
@@ -1,3 +1,4 @@
+import concurrent.futures
 import copy
 import itertools
 import multiprocessing
@@ -171,6 +172,17 @@ def test_optimize_parallel(n_trials: int, n_jobs: int, storage_mode: str) -> Non
         check_study(study)
 
 
+def test_optimize_with_thread_pool_executor() -> None:
+    def objective(t: Trial) -> float:
+        return t.suggest_float("x", -10, 10)
+
+    study = create_study()
+    with concurrent.futures.ThreadPoolExecutor(max_workers=5) as pool:
+        for _ in range(10):
+            pool.submit(study.optimize, objective, n_trials=10)
+    assert len(study.trials) == 100
+
+
 @pytest.mark.parametrize(
     "n_trials, n_jobs, storage_mode",
     itertools.product(
@@ -248,6 +260,21 @@ def test_optimize_with_reseeding(n_jobs: int, storage_mode: str) -> None:
             assert mock_object.call_count == 1
 
 
+def test_call_another_study_optimize_in_optimize() -> None:
+    def inner_objective(t: Trial) -> float:
+        return t.suggest_float("x", -10, 10)
+
+    def objective(t: Trial) -> float:
+        inner_study = create_study()
+        inner_study.enqueue_trial({"x": t.suggest_int("initial_point", -10, 10)})
+        inner_study.optimize(inner_objective, n_trials=10)
+        return inner_study.best_value
+
+    study = create_study()
+    study.optimize(objective, n_trials=10)
+    assert len(study.trials) == 10
+
+
 @pytest.mark.parametrize("storage_mode", STORAGE_MODES)
 def test_study_set_and_get_user_attrs(storage_mode: str) -> None: