Repository: optuna/optuna
Pull request: Resolve the issue where `BruteForceSampler` fails to suggest all combinations
<!-- 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
<!-- Describe your motivation why you will submit this PR. This is useful for reviewers to understand the context of PR. -->

Fixes #5780 

## Description of the changes
<!-- Describe the changes in this PR. -->

When running `BruteForceSampler` in parallel, there was a bug where a node being processed by one thread was forcibly changed to a leaf-state. This caused other threads to mistakenly assume there were no further nodes to explore beyond that point, leading to premature termination before all combinations were explored.

To address this issue, this PR introduces an `is_running` flag to the nodes. This flag allows distinguishing between nodes that are simply unexpanded and those that are unexpanded because they are currently being processed.


Problem statement:
`BruteForceSampler` stops before suggesting all combinations in parallel optimization
### Expected behavior

`BruteForceSampler` should suggest all parameter combinations, but it stops unexpectedly before that in parallel optimization.

### Environment

- Optuna version:4.2.0.dev
- Python version:3.9.4
- OS:macOS-15.1-x86_64-i386-64bit
- (Optional) Other libraries and their versions:


### Error messages, stack traces, or logs

```shell
See the following MRE.
```


### Steps to reproduce

```python
import time
import optuna

def objective(trial):
    x = trial.suggest_int("x", 0, 1)
    time.sleep(x)
    y = trial.suggest_int("y", 0, 1)
    return x + y

sampler = optuna.samplers.BruteForceSampler(seed=42)
study = optuna.create_study(sampler=sampler)
study.optimize(objective, n_jobs=2)
```

```
[I 2024-11-22 17:43:20,854] A new study created in memory with name: no-name-86ca63d9-e334-4237-a5a5-854dd2483f9f
[I 2024-11-22 17:43:20,855] Trial 0 finished with value: 1.0 and parameters: {'x': 0, 'y': 1}. Best is trial 0 with value: 1.0.
[I 2024-11-22 17:43:20,856] Trial 2 finished with value: 0.0 and parameters: {'x': 0, 'y': 0}. Best is trial 2 with value: 0.0.
[I 2024-11-22 17:43:21,861] Trial 1 finished with value: 1.0 and parameters: {'x': 1, 'y': 0}. Best is trial 2 with value: 0.0.
```


### Additional context (optional)

_No response_

Hints from issue discussion:

Commit message: Update test_parallel_optimize_with_sleep

Patch to review:
diff --git a/optuna/samplers/_brute_force.py b/optuna/samplers/_brute_force.py
index 58a4022f851..50ce4ba8ab3 100644
--- a/optuna/samplers/_brute_force.py
+++ b/optuna/samplers/_brute_force.py
@@ -36,6 +36,7 @@ class _TreeNode:
 
     param_name: str | None = None
     children: dict[float, "_TreeNode"] | None = None
+    is_running: bool = False
 
     def expand(self, param_name: str | None, search_space: Iterable[float]) -> None:
         # If the node is unexpanded, expand it.
@@ -52,6 +53,9 @@ def expand(self, param_name: str | None, search_space: Iterable[float]) -> None:
                     f"search_space mismatch: {set(self.children.keys())} != {set(search_space)}"
                 )
 
+    def set_running(self) -> None:
+        self.is_running = True
+
     def set_leaf(self) -> None:
         self.expand(None, [])
 
@@ -68,23 +72,38 @@ def add_path(
             current_node = current_node.children[value]
         return current_node
 
-    def count_unexpanded(self) -> int:
+    def count_unexpanded(self, exclude_running: bool) -> int:
         # Count the number of unexpanded nodes in the subtree.
         return (
             1
-            if self.children is None
-            else sum(child.count_unexpanded() for child in self.children.values())
+            if self.children is None and (not exclude_running or not self.is_running)
+            else (
+                0
+                if self.children is None
+                else sum(
+                    child.count_unexpanded(exclude_running) for child in self.children.values()
+                )
+            )
         )
 
-    def sample_child(self, rng: np.random.RandomState) -> float:
+    def sample_child(self, rng: np.random.RandomState, exclude_running: bool) -> float:
         assert self.children is not None
         # Sample an unexpanded node in the subtree uniformly, and return the first
         # parameter value in the path to the node.
         # Equivalently, we sample the child node with weights proportional to the number
         # of unexpanded nodes in the subtree.
         weights = np.array(
-            [child.count_unexpanded() for child in self.children.values()], dtype=np.float64
+            [child.count_unexpanded(exclude_running) for child in self.children.values()],
+            dtype=np.float64,
         )
+        if any(
+            not value.is_running and weights[i] > 0
+            for i, value in enumerate(self.children.values())
+        ):
+            # Prioritize picking non-running and unexpanded nodes.
+            for i, child in enumerate(self.children.values()):
+                if child.is_running:
+                    weights[i] = 0.0
         weights /= weights.sum()
         return rng.choice(list(self.children.keys()), p=weights)
 
@@ -130,10 +149,14 @@ def objective(trial):
             that it is not recommended using this option in distributed optimization settings since
             this option cannot ensure the order of trials and may increase the number of duplicate
             suggestions during distributed optimization.
+        enforce_unique_samples:
+            If this is :obj:`True`, the sampler strictly avoids duplicate samples. Please note that
+            the sampler may stop before trying the entire search space when this option is enabled.
     """
 
-    def __init__(self, seed: int | None = None) -> None:
+    def __init__(self, seed: int | None = None, enforce_unique_samples: bool = False) -> None:
         self._rng = LazyRandomState(seed)
+        self._enforce_unique_samples = enforce_unique_samples
 
     def infer_relative_search_space(
         self, study: Study, trial: FrozenTrial
@@ -150,7 +173,6 @@ 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()):
                 continue
@@ -170,12 +192,7 @@ def _populate_tree(
                 if trial.state.is_finished():
                     leaf.set_leaf()
                 else:
-                    incomplete_leaves.append(leaf)
-
-        # 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()
+                    leaf.set_running()
 
     def sample_independent(
         self,
@@ -184,6 +201,8 @@ def sample_independent(
         param_name: str,
         param_distribution: BaseDistribution,
     ) -> Any:
+        exclude_running = self._enforce_unique_samples
+
         trials = study.get_trials(
             deepcopy=False,
             states=(
@@ -200,10 +219,12 @@ def sample_independent(
         # being initialized as an empty graph, which is created with n_jobs > 1
         # where we get trials[i].params = {} for some i.
         self._populate_tree(tree, (t for t in trials if t.number != trial.number), trial.params)
-        if tree.count_unexpanded() == 0:
+        if tree.count_unexpanded(exclude_running) == 0:
             return param_distribution.to_external_repr(self._rng.rng.choice(candidates))
         else:
-            return param_distribution.to_external_repr(tree.sample_child(self._rng.rng))
+            return param_distribution.to_external_repr(
+                tree.sample_child(self._rng.rng, exclude_running)
+            )
 
     def after_trial(
         self,
@@ -212,6 +233,8 @@ def after_trial(
         state: TrialState,
         values: Sequence[float] | None,
     ) -> None:
+        exclude_running = self._enforce_unique_samples
+
         trials = study.get_trials(
             deepcopy=False,
             states=(
@@ -240,7 +263,7 @@ def after_trial(
             {},
         )
 
-        if tree.count_unexpanded() == 0:
+        if tree.count_unexpanded(exclude_running) == 0:
             study.stop()
 
 
diff --git a/tests/samplers_tests/test_brute_force.py b/tests/samplers_tests/test_brute_force.py
index 2de6f95cdaa..01b74cdb716 100644
--- a/tests/samplers_tests/test_brute_force.py
+++ b/tests/samplers_tests/test_brute_force.py
@@ -1,3 +1,5 @@
+import time
+
 import numpy as np
 import pytest
 
@@ -88,7 +90,7 @@ def test_tree_node_count_unexpanded() -> None:
             2: _TreeNode(),
         },
     )
-    assert tree.count_unexpanded() == 3
+    assert tree.count_unexpanded(exclude_running=False) == 3
 
 
 def test_study_optimize_with_single_search_space() -> None:
@@ -299,3 +301,21 @@ def test_parallel_optimize() -> None:
     x1 = trial1.suggest_categorical("x", ["a", "b"])
     x2 = trial2.suggest_categorical("x", ["a", "b"])
     assert {x1, x2} == {"a", "b"}
+
+
+def test_parallel_optimize_with_sleep() -> None:
+    def objective(trial: Trial) -> float:
+        x = trial.suggest_int("x", 0, 1)
+        time.sleep(x)
+        y = trial.suggest_int("y", 0, 1)
+        return x + y
+
+    # Seed is fixed to reproduce the same result.
+    # See: https://github.com/optuna/optuna/issues/5780
+    study = optuna.create_study(sampler=samplers.BruteForceSampler(seed=42))
+    study.optimize(objective, n_jobs=2)
+
+    expected_suggested_values = [{"x": i, "y": j} for i in range(2) for j in range(2)]
+    all_suggested_values = [t.params for t in study.trials]
+    for expected_value in expected_suggested_values:
+        assert expected_value in all_suggested_values


Code review comments:
diff_hunk: @@ -88,7 +90,7 @@ def test_tree_node_count_unexpanded() -> None:
             2: _TreeNode(),
         },
     )
-    assert tree.count_unexpanded() == 3
+    assert tree.count_unexpanded(exclude_running=False) == 3
 
path: tests/samplers_tests/test_brute_force.py
text: @user1:
Could you change this test to appropriately set `is_running` so that `exclude_running=False` and `exclude_running=True` returns different results, and assert both cases?

diff_hunk: @@ -130,10 +149,14 @@ def objective(trial):
             that it is not recommended using this option in distributed optimization settings since
             this option cannot ensure the order of trials and may increase the number of duplicate
             suggestions during distributed optimization.
+        enforce_unique_samples:
+            If this is :obj:`True`, the sampler strictly avoids duplicate samples. Please note that
+            the sampler may stop before trying the entire search space when this option is enabled.
     """
 
-    def __init__(self, seed: int | None = None) -> None:
+    def __init__(self, seed: int | None = None, enforce_unique_samples: bool = False) -> None:
path: optuna/samplers/_brute_force.py
text: @user1:
Setting `exclude_running = True` does not strictly avoid duplicate samples because we do not hold any lock on the storage (two workers may sample the same parameter at the same time and we have no way to strictly prevent that). So, a different name is preferred. How about "avoid_premature_stop" (the boolean needs to be flipped)? 

@user1:
Also, since the original implementation was `enforce_unique_samples = True`, IMO the default should be so to prevent breaking changes (users may observe that working code suddenly starts producing duplicate trials and e.g. wasting money on AWS). What do you think?

diff_hunk: @@ -68,23 +72,38 @@ def add_path(
             current_node = current_node.children[value]
         return current_node
 
-    def count_unexpanded(self) -> int:
+    def count_unexpanded(self, exclude_running: bool) -> int:
         # Count the number of unexpanded nodes in the subtree.
         return (
             1
-            if self.children is None
-            else sum(child.count_unexpanded() for child in self.children.values())
+            if self.children is None and (not exclude_running or not self.is_running)
+            else (
+                0
+                if self.children is None
+                else sum(
+                    child.count_unexpanded(exclude_running) for child in self.children.values()
+                )
+            )
         )
path: optuna/samplers/_brute_force.py
text: @user1:
A better way to implement this function is 
```
        # max(0, ...) is taken to prevent negative values when a node happens to be running on a fully-expanded subtree.
        return max(0, (
            1
            if self.children is None
            else sum(child.count_unexpanded(exclude_running) for child in self.children.values())
        ) - int(self.is_running and exclude_running))
```
This would be advantageous when there is a running trial on already-expanded node.
        

@user1:
Sorry, I noticed that this may cause premature stop in additional cases.