Repository: optuna/optuna

Issue:
`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_

Discussion:

Solution patch:
diff --git a/README.md b/README.md
index 9527cbc28..023ecc30c 100644
--- a/README.md
+++ b/README.md
@@ -25,12 +25,12 @@ Optuna can dynamically construct the search spaces for the hyperparameters.
 
 ## :loudspeaker: News
 <!-- TODO: when you add a new line, please delete the oldest line -->
+* **Jan 22, 2025**: A new article [OptunaHub Benchmarks: A New Feature to Use/Register Various Benchmark Problems](https://medium.com/optuna/optunahub-benchmarks-a-new-feature-to-use-register-various-benchmark-problems-694401524ce0) has been published.
+* **Jan 20, 2025**: Optuna 4.2.0 and OptunaHub 0.2.0 are out! Try the newest Optuna and OptunaHub! Check out [the release note](https://github.com/optuna/optuna/releases/tag/v4.2.0) for details.
+* **Jan 16, 2025**: A new article [Overview of Python Free Threading (v3.13t) Support in Optuna](https://medium.com/optuna/overview-of-python-free-threading-v3-13t-support-in-optuna-ad9ab62a11ba) has been published.
 * **Nov 12, 2024**: We released Optuna 4.1 with new features, Python 3.13 support and much more! Check out [the release note](https://github.com/optuna/optuna/releases/tag/v4.1.0) for details.
 * **Nov 6, 2024**: A new article [AutoSampler: Automatic Selection of Optimization Algorithms in Optuna](https://medium.com/optuna/autosampler-automatic-selection-of-optimization-algorithms-in-optuna-1443875fd8f9) has been published.
 * **Oct 21, 2024**: We posted [an article](https://medium.com/optuna/an-introduction-to-moea-d-and-examples-of-multi-objective-optimization-comparisons-8630565a4e89) introducing [MOEA/D](https://hub.optuna.org/samplers/moead/) and an example comparison with other optimization methods.
-* **Oct 15, 2024**: We posted [an article](https://medium.com/optuna/introducing-a-new-terminator-early-termination-of-black-box-optimization-based-on-expected-9a660774fcdb) about `Terminator`, which is expanded in Optuna 4.0.
-* **Sep 18, 2024**: We posted [an article](https://medium.com/optuna/introducing-the-stabilized-journalstorage-in-optuna-4-0-from-mechanism-to-use-case-e320795ffb61) about `JournalStorage`, which is stabilized in Optuna 4.0.
-* **Sep 2, 2024**: Optuna 4.0 is available!  You can install it by `pip install -U optuna`. Find the latest [here](https://github.com/optuna/optuna/releases) and check [our article](https://medium.com/optuna/optuna-4-0-whats-new-in-the-major-release-3325a8420d10).
 
 ## :fire: Key Features
 
diff --git a/optuna/samplers/_brute_force.py b/optuna/samplers/_brute_force.py
index 58a4022f8..ee07cedac 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 @@ class _TreeNode:
                     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,31 @@ class _TreeNode:
             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:
+            return 0 if exclude_running and self.is_running else 1
+        else:
+            return 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 +142,17 @@ class BruteForceSampler(BaseSampler):
             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.
+        avoid_premature_stop:
+            If :obj:`True`, the sampler performs a strict exhaustive search. Please note
+            that enabling this option may increase the likelihood of duplicate sampling.
+            When this option is not enabled (default), the sampler applies a looser criterion for
+            determining when to stop the search, which may result in incomplete coverage of the
+            search space. For more information, see https://github.com/optuna/optuna/issues/5780.
     """
 
-    def __init__(self, seed: int | None = None) -> None:
+    def __init__(self, seed: int | None = None, avoid_premature_stop: bool = False) -> None:
         self._rng = LazyRandomState(seed)
+        self._avoid_premature_stop = avoid_premature_stop
 
     def infer_relative_search_space(
         self, study: Study, trial: FrozenTrial
@@ -150,7 +169,6 @@ class BruteForceSampler(BaseSampler):
         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 +188,7 @@ class BruteForceSampler(BaseSampler):
                 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 +197,8 @@ class BruteForceSampler(BaseSampler):
         param_name: str,
         param_distribution: BaseDistribution,
     ) -> Any:
+        exclude_running = not self._avoid_premature_stop
+
         trials = study.get_trials(
             deepcopy=False,
             states=(
@@ -200,10 +215,12 @@ class BruteForceSampler(BaseSampler):
         # 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 +229,8 @@ class BruteForceSampler(BaseSampler):
         state: TrialState,
         values: Sequence[float] | None,
     ) -> None:
+        exclude_running = not self._avoid_premature_stop
+
         trials = study.get_trials(
             deepcopy=False,
             states=(
@@ -240,7 +259,7 @@ class BruteForceSampler(BaseSampler):
             {},
         )
 
-        if tree.count_unexpanded() == 0:
+        if tree.count_unexpanded(exclude_running) == 0:
             study.stop()
 
 
diff --git a/optuna/testing/storages.py b/optuna/testing/storages.py
index 7a98ac0c4..a267c67ef 100644
--- a/optuna/testing/storages.py
+++ b/optuna/testing/storages.py
@@ -2,6 +2,7 @@ from __future__ import annotations
 
 import socket
 import threading
+import time
 from types import TracebackType
 from typing import Any
 from typing import IO
@@ -98,7 +99,16 @@ class StorageSupplier:
             self.thread = threading.Thread(target=self.server.start)
             self.thread.start()
 
-            return GrpcStorageProxy(host="localhost", port=port)
+            proxy = GrpcStorageProxy(host="localhost", port=port)
+
+            # Wait until the server is ready.
+            while True:
+                try:
+                    proxy.get_all_studies()
+                    return proxy
+                except grpc.RpcError:
+                    time.sleep(1)
+                    continue
         else:
             assert False
 
diff --git a/optuna/version.py b/optuna/version.py
index 1e417d3e3..fc4c5571e 100644
--- a/optuna/version.py
+++ b/optuna/version.py
@@ -1 +1 @@
-__version__ = "4.2.0.dev"
+__version__ = "4.3.0.dev"
diff --git a/optuna/visualization/matplotlib/_contour.py b/optuna/visualization/matplotlib/_contour.py
index 6a4737d9d..5134a3797 100644
--- a/optuna/visualization/matplotlib/_contour.py
+++ b/optuna/visualization/matplotlib/_contour.py
@@ -175,7 +175,12 @@ def _calculate_axis_data(
     returned_values: Sequence[int | float]
     if axis.is_cat:
         enc = _LabelEncoder()
-        returned_values = enc.fit_transform(list(map(str, values)))
+        # Fit LabelEncoder with all the categories in categorical distribution.
+        enc.fit(list(map(str, filter(lambda value: value is not None, axis.values))))
+        # Then transform the values using the fitted label encoder.
+        # Note that `values` may not include all the categories,
+        # so we use `axis.values` for fitting.
+        returned_values = enc.transform(list(map(str, values)))
         cat_param_labels = enc.get_labels()
         cat_param_pos = enc.get_indices()
     else:


Test patch:
diff --git a/tests/samplers_tests/test_brute_force.py b/tests/samplers_tests/test_brute_force.py
index 2de6f95cd..745d306ab 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
 
@@ -85,10 +87,11 @@ def test_tree_node_count_unexpanded() -> None:
                     1.0: _TreeNode(),
                 },
             ),
-            2: _TreeNode(),
+            2: _TreeNode(is_running=True),
         },
     )
-    assert tree.count_unexpanded() == 3
+    assert tree.count_unexpanded(exclude_running=False) == 3
+    assert tree.count_unexpanded(exclude_running=True) == 2
 
 
 def test_study_optimize_with_single_search_space() -> None:
@@ -299,3 +302,39 @@ 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": 0, "y": 0},
+        {"x": 0, "y": 1},
+        {"x": 1, "y": 0},
+    ]
+    all_suggested_values = [t.params for t in study.trials]
+    assert len(all_suggested_values) == len(expected_suggested_values)
+    for expected_value in expected_suggested_values:
+        assert expected_value in all_suggested_values
+
+    study = optuna.create_study(
+        sampler=samplers.BruteForceSampler(seed=42, avoid_premature_stop=True)
+    )
+    study.optimize(objective, n_jobs=2)
+    expected_suggested_values = [
+        {"x": 0, "y": 0},
+        {"x": 0, "y": 1},
+        {"x": 1, "y": 0},
+        {"x": 1, "y": 1},
+    ]
+    all_suggested_values = [t.params for t in study.trials]
+    for expected_value in expected_suggested_values:
+        assert expected_value in all_suggested_values