Repository: optuna/optuna
Pull request: Fix `BruteForceSampler` in parallel optimization
<!-- Thank you for creating a pull request! In general, we merge your pull request after it gets two or more approvals. To proceed to the review process by the maintainers, please make sure that the PR meets the following conditions: (1) it passes all CI checks, and (2) it is neither in the draft nor WIP state. If you wish to discuss the PR in the draft state or need any other help, please mention the Optuna development team in the PR. -->

## Motivation
Fix #4948.
<!-- Describe your motivation why you will submit this PR. This is useful for reviewers to understand the context of PR. -->

## Description of the changes
https://github.com/optuna/optuna/issues/4948#issuecomment-1756861939 explains why this bug happens. This PR takes the current trial into account when inferring whether each trial is "finalized". 


Problem statement:
`BruteForceSampler` crashes when used with `n_jobs`
### Expected behavior

`BruteForceSampler` should work with `n_jobs`, but it can raise a `ValueError`.

### Environment

- Optuna version:3.4.0.dev
- Python version:3.8.6
- OS:macOS-10.16-x86_64-i386-64bit
- (Optional) Other libraries and their versions:


### Error messages, stack traces, or logs

```shell
$ python grid.py 
grid.py:8: ExperimentalWarning: BruteForceSampler is experimental (supported from v3.1.0). The interface can change in the future.
  sampler = optuna.samplers.BruteForceSampler(seed=42)
[I 2023-09-26 10:37:21,476] A new study created in RDB with name: no-name-1a1d7436-f6cb-4c45-b3e5-b2c4ab137846
[W 2023-09-26 10:37:22,585] Trial 0 failed with parameters: {} because of the following error: ValueError('param_name mismatch: None != x').
Traceback (most recent call last):
  File "/Users/naotomizuno/optuna/optuna/study/_optimize.py", line 200, in _run_trial
    value_or_values = func(trial)
  File "grid.py", line 6, in objective
    return trial.suggest_int("x", 0, 9)
  File "/Users/naotomizuno/optuna/optuna/trial/_trial.py", line 321, in suggest_int
    suggested_value = int(self._suggest(name, distribution))
  File "/Users/naotomizuno/optuna/optuna/trial/_trial.py", line 630, in _suggest
    param_value = self.study.sampler.sample_independent(
  File "/Users/naotomizuno/optuna/optuna/samplers/_brute_force.py", line 195, in sample_independent
    tree.expand(param_name, candidates)
  File "/Users/naotomizuno/optuna/optuna/samplers/_brute_force.py", line 46, in expand
    raise ValueError(f"param_name mismatch: {self.param_name} != {param_name}")
ValueError: param_name mismatch: None != x
[W 2023-09-26 10:37:22,587] Trial 0 failed with value None.
[W 2023-09-26 10:37:22,593] Trial 1 failed with parameters: {} because of the following error: ValueError('param_name mismatch: None != x').
Traceback (most recent call last):
  File "/Users/naotomizuno/optuna/optuna/study/_optimize.py", line 200, in _run_trial
    value_or_values = func(trial)
  File "grid.py", line 6, in objective
    return trial.suggest_int("x", 0, 9)
  File "/Users/naotomizuno/optuna/optuna/trial/_trial.py", line 321, in suggest_int
    suggested_value = int(self._suggest(name, distribution))
  File "/Users/naotomizuno/optuna/optuna/trial/_trial.py", line 630, in _suggest
    param_value = self.study.sampler.sample_independent(
  File "/Users/naotomizuno/optuna/optuna/samplers/_brute_force.py", line 195, in sample_independent
    tree.expand(param_name, candidates)
  File "/Users/naotomizuno/optuna/optuna/samplers/_brute_force.py", line 46, in expand
    raise ValueError(f"param_name mismatch: {self.param_name} != {param_name}")
ValueError: param_name mismatch: None != x
[W 2023-09-26 10:37:22,593] Trial 1 failed with value None.
Traceback (most recent call last):
  File "grid.py", line 11, in <module>
    study.optimize(objective, n_jobs=2)
  File "/Users/naotomizuno/optuna/optuna/study/study.py", line 439, in optimize
    _optimize(
  File "/Users/naotomizuno/optuna/optuna/study/_optimize.py", line 103, in _optimize
    f.result()
  File "/Users/naotomizuno/.pyenv/versions/3.8.6/lib/python3.8/concurrent/futures/_base.py", line 432, in result
    return self.__get_result()
  File "/Users/naotomizuno/.pyenv/versions/3.8.6/lib/python3.8/concurrent/futures/_base.py", line 388, in __get_result
    raise self._exception
  File "/Users/naotomizuno/.pyenv/versions/3.8.6/lib/python3.8/concurrent/futures/thread.py", line 57, in run
    result = self.fn(*self.args, **self.kwargs)
  File "/Users/naotomizuno/optuna/optuna/study/_optimize.py", line 163, in _optimize_sequential
    frozen_trial = _run_trial(study, func, catch)
  File "/Users/naotomizuno/optuna/optuna/study/_optimize.py", line 251, in _run_trial
    raise func_err
  File "/Users/naotomizuno/optuna/optuna/study/_optimize.py", line 200, in _run_trial
    value_or_values = func(trial)
  File "grid.py", line 6, in objective
    return trial.suggest_int("x", 0, 9)
  File "/Users/naotomizuno/optuna/optuna/trial/_trial.py", line 321, in suggest_int
    suggested_value = int(self._suggest(name, distribution))
  File "/Users/naotomizuno/optuna/optuna/trial/_trial.py", line 630, in _suggest
    param_value = self.study.sampler.sample_independent(
  File "/Users/naotomizuno/optuna/optuna/samplers/_brute_force.py", line 195, in sample_independent
    tree.expand(param_name, candidates)
  File "/Users/naotomizuno/optuna/optuna/samplers/_brute_force.py", line 46, in expand
    raise ValueError(f"param_name mismatch: {self.param_name} != {param_name}")
ValueError: param_name mismatch: None != x
```


### Steps to reproduce

Run the following code.

```python
import time
import optuna

def objective(trial):
    time.sleep(1)
    return trial.suggest_int("x", 0, 9)

sampler = optuna.samplers.BruteForceSampler(seed=42)
storage = "sqlite:///tmp.db"
study = optuna.create_study(storage=storage, sampler=sampler)
study.optimize(objective, n_jobs=2)

```


### Additional context (optional)

_No response_

Hints from issue discussion:
It looks like this issue is raised by the following code:

Pull Request: https://github.com/optuna/optuna/pull/5008

```
# Add all incomplete leaf nodes at the end because they may not have complete search space.
for leaf in incomplete_leaves:
    if leaf.children is None:
        leaf.set_leaf()
```

A node is set as a leaf node if the state of a trial is finished, while the above code sets all incomplete leaves as leaf nodes.

```
# The parameters are on the defined grid.
if trial.state.is_finished():
    leaf.set_leaf()
else:
    incomplete_leaves.append(leaf)
```
Thank you for the bug report. The bug is likely caused by the following reason:
BruteForceSampler builds a parameter tree in each call to `sample_independent` to track which parameter combinations are unexplored. Since Optuna doesn't have an interface to "finalize" parameters before a trial is finished, we assume that running trials have fixed parameters _unless there are another trial that has superset of parameters_. When two new trials (with no parameters sampled yet) are running, this logic assumes that the objective function has zero parameters, and thus raises an error when a parameter is suggested.
Failing test:
```

def test_parallel_optimize() -> None:
    study = optuna.create_study(sampler=samplers.BruteForceSampler())
    trial1 = study.ask()
    trial2 = study.ask()
    x1 = trial1.suggest_categorical("x", ["a", "b"])
    x2 = trial2.suggest_categorical("x", ["a", "b"])
    assert {x1, x2} == {"a", "b"}
```

Commit message: Update optuna/samplers/_brute_force.py

Co-authored-by: Shuhei Watanabe <47781922+nabenabe0928@users.noreply.github.com>

Patch to review:
diff --git a/optuna/samplers/_brute_force.py b/optuna/samplers/_brute_force.py
index a08c33c224c..a9fd5442a82 100644
--- a/optuna/samplers/_brute_force.py
+++ b/optuna/samplers/_brute_force.py
@@ -143,9 +143,10 @@ def sample_relative(
         return {}
 
     @staticmethod
-    def _build_tree(trials: Iterable[FrozenTrial], params: Dict[str, Any]) -> _TreeNode:
-        # Build a _TreeNode under given params from the given trials.
-        tree = _TreeNode()
+    def _populate_tree(
+        tree: _TreeNode, trials: Iterable[FrozenTrial], params: Dict[str, Any]
+    ) -> None:
+        # Populate tree under given params from the given trials.
         incomplete_leaves: List[_TreeNode] = []
         for trial in trials:
             if not all(p in trial.params and trial.params[p] == v for p, v in params.items()):
@@ -172,7 +173,6 @@ def _build_tree(trials: Iterable[FrozenTrial], params: Dict[str, Any]) -> _TreeN
         for leaf in incomplete_leaves:
             if leaf.children is None:
                 leaf.set_leaf()
-        return tree
 
     def sample_independent(
         self,
@@ -190,9 +190,13 @@ def sample_independent(
                 TrialState.FAIL,
             ),
         )
-        tree = self._build_tree((t for t in trials if t.number != trial.number), trial.params)
+        tree = _TreeNode()
         candidates = _enumerate_candidates(param_distribution)
         tree.expand(param_name, candidates)
+        # Populating must happen after the initialization above to prevent tree from
+        # being initialized as an empty graph, which is created n_jobs > 1
+        # where we get trials.params = {}.
+        self._populate_tree(tree, (t for t in trials if t.number != trial.number), trial.params)
         if tree.count_unexpanded() == 0:
             return param_distribution.to_external_repr(self._rng.choice(candidates))
         else:
@@ -214,7 +218,9 @@ def after_trial(
                 TrialState.FAIL,
             ),
         )
-        tree = self._build_tree(
+        tree = _TreeNode()
+        self._populate_tree(
+            tree,
             (
                 t
                 if t.number != trial.number
diff --git a/tests/samplers_tests/test_brute_force.py b/tests/samplers_tests/test_brute_force.py
index 3ae06ced8e5..5eee437b5fb 100644
--- a/tests/samplers_tests/test_brute_force.py
+++ b/tests/samplers_tests/test_brute_force.py
@@ -290,3 +290,12 @@ def objective(trial: Trial) -> float:
     assert len(all_suggested_values) == len(expected_suggested_values)
     for a in expected_suggested_values:
         assert a in all_suggested_values
+
+
+def test_parallel_optimize() -> None:
+    study = optuna.create_study(sampler=samplers.BruteForceSampler())
+    trial1 = study.ask()
+    trial2 = study.ask()
+    x1 = trial1.suggest_categorical("x", ["a", "b"])
+    x2 = trial2.suggest_categorical("x", ["a", "b"])
+    assert {x1, x2} == {"a", "b"}


Code review comments:
diff_hunk: @@ -190,9 +190,13 @@ def sample_independent(
                 TrialState.FAIL,
             ),
         )
-        tree = self._build_tree((t for t in trials if t.number != trial.number), trial.params)
+        tree = _TreeNode()
         candidates = _enumerate_candidates(param_distribution)
         tree.expand(param_name, candidates)
+        # Populating must happen after the initialization above to prevent tree from
+        # being initialized as an empty graph, which is created n_jobs > 1
+        # where we get trials.params = {}.
path: optuna/samplers/_brute_force.py
text: @user1:
Sorry, I noticed some mistakes in my suggestion.

```suggestion
        # Populating must happen after the initialization above to prevent `tree` from
        # being initialized as an empty graph, which is created with n_jobs > 1
        # where we get trials[i].params = {} for some i.
```