Repository: optuna/optuna
Pull request: Fix bug of nondeterministic behavior of TPESampler when `group=True`
## Motivation
Fix #3137

Due to Python's nondeterministic string hashing, the iteration of `Dict[str, BaseDistribution]` is in nondeterministic order.

## Description of the changes
Sorted the keys to make the iteration deterministic.



I don't think that this PR is perfect to resolve the original issue because:
* It is very hard for users and developers to always keep in mind that iterating distributions is nondeterministic while development.
* Moreover, such situation may usually happen because `Dict[str, BaseDistribution]` appears in various positions in the codebase.

One of possible solutions may be creating `SearchSpace` class instead of the raw `Dict[str, BaseDistribution]` and manually handle hashing, but this may require a large-scale modification.

If you have any ideas for modification, please let me know.

Problem statement:
Non-determinism in TPE sampler when `group=True`
TPE sampler is non-deterministic in different runs (with the same seed, no parallelization) when group decomposed search space is enabled.

## Expected behavior

Consider the script below:
```python

import optuna

sampler = optuna.samplers.TPESampler(n_startup_trials=1, seed=2, multivariate=True, group=True)
study = optuna.create_study(sampler=sampler)

for _ in range(10):
    trial = study.ask()
    picked = [i for i in range(10) if trial.suggest_int(str(i), 0, 1) == 1]
    print(picked)
    study.tell(trial, len(picked))

```

Different runs of this script are expected to yield the same output (because the seed is fixed, the objective function is deterministic, and we are running in sequential mode), but they do not. For example, two consecutive runs of this script yielded the following outputs (your output might differ) :

First run:
```
[2, 7]
[2, 7]
[2, 7]
[5, 6, 7, 8, 9]
[1, 5]
[0, 2, 3, 4, 7]
[2, 3, 8]
[1, 2, 3, 6, 7]
[0, 1, 2, 4]
[3, 5, 7]
```

Second run:
```
[2, 7]
[2, 7, 9]
[2, 7]
[1, 6]
[2, 3, 6, 7, 8]
[0, 2, 3, 5, 7]
[1, 7, 8]
[0, 1, 2, 6, 7]
[3]
[3, 5]
```

Note that if we set `group=False` when initializing the TPE sampler, we get deterministic outputs.

## Environment

- Optuna version: 2.91
- Python version: 3.7.4
- OS: Linux

## Error messages, stack traces, or logs

Not applicable.

## Steps to reproduce

Run the script above multiple times. Observe that the output changes everytime.

Hints from issue discussion:
Thank you for pointing it out! I found this bug caused by the randomness of Python's hash seed. This line is not deterministic. https://github.com/optuna/optuna/blob/2a9f8639c1703f455fdbad211e7fa8ab53628a21/optuna/samplers/_search_space/group_decomposed.py#L21

Commit message: Add test

Patch to review:
diff --git a/optuna/samplers/_tpe/sampler.py b/optuna/samplers/_tpe/sampler.py
index 94ee0873e45..f1393f4f87d 100644
--- a/optuna/samplers/_tpe/sampler.py
+++ b/optuna/samplers/_tpe/sampler.py
@@ -303,7 +303,8 @@ def infer_relative_search_space(
             assert self._group_decomposed_search_space is not None
             self._search_space_group = self._group_decomposed_search_space.calculate(study)
             for sub_space in self._search_space_group.search_spaces:
-                for name, distribution in sub_space.items():
+                # Sort keys because Python's string hashing is nondeterministic.
+                for name, distribution in sorted(sub_space.items()):
                     if distribution.single():
                         continue
                     search_space[name] = distribution
@@ -340,7 +341,8 @@ def sample_relative(
             params = {}
             for sub_space in self._search_space_group.search_spaces:
                 search_space = {}
-                for name, distribution in sub_space.items():
+                # Sort keys because Python's string hashing is nondeterministic.
+                for name, distribution in sorted(sub_space.items()):
                     if not distribution.single():
                         search_space[name] = distribution
                 params.update(self._sample_relative(study, trial, search_space))
diff --git a/tests/samplers_tests/tpe_tests/test_sampler.py b/tests/samplers_tests/tpe_tests/test_sampler.py
index 8ee2164c139..7575979007f 100644
--- a/tests/samplers_tests/tpe_tests/test_sampler.py
+++ b/tests/samplers_tests/tpe_tests/test_sampler.py
@@ -1,6 +1,7 @@
 import random
 from typing import Callable
 from typing import Dict
+from typing import List
 from typing import Optional
 from typing import Union
 from unittest.mock import Mock
@@ -16,6 +17,7 @@
 from optuna import TrialPruned
 from optuna.samplers import _tpe
 from optuna.samplers import TPESampler
+from optuna.study.study import create_study
 from optuna.trial import Trial
 
 
@@ -1054,6 +1056,37 @@ def test_group_experimental_warning() -> None:
         _ = TPESampler(multivariate=True, group=True)
 
 
+def run_tpe(k: int, sequence_dict: Dict[int, List[int]], hash_dict: Dict[int, int]) -> None:
+    hash_dict[k] = hash("nondeterministic hash")
+    sampler = TPESampler(n_startup_trials=1, seed=2, multivariate=True, group=True)
+    study = create_study(sampler=sampler)
+    sequence = []
+    for _ in range(10):
+        trial = study.ask()
+        picked = [i for i in range(10) if trial.suggest_int(str(i), 0, 1) == 1]
+        study.tell(trial, len(picked))
+        sequence.extend(picked)
+    sequence_dict[k] = sequence
+
+
+# Test deterministic iteration as discussed in https://github.com/optuna/optuna/issues/3137
+def test_group_deterministic_iteration() -> None:
+    from multiprocessing import Manager
+    from multiprocessing import Process
+
+    manager = Manager()
+    sequence_dict: Dict[int, List[int]] = manager.dict()
+    hash_dict: Dict[int, int] = manager.dict()
+    for i in range(3):
+        p = Process(target=run_tpe, args=(i, sequence_dict, hash_dict))
+        p.start()
+        p.join()
+    # Hashes are expected to be different because string hashing is nondeterministic per process.
+    assert not (hash_dict[0] == hash_dict[1] == hash_dict[2])
+    # But the sequences are expected to be the same.
+    assert sequence_dict[0] == sequence_dict[1] == sequence_dict[2]
+
+
 @pytest.mark.parametrize("direction", ["minimize", "maximize"])
 @pytest.mark.parametrize("multivariate", [True, False])
 def test_constant_liar_observation_pairs(direction: str, multivariate: bool) -> None:


Code review comments:
diff_hunk: @@ -1054,6 +1056,37 @@ def test_group_experimental_warning() -> None:
         _ = TPESampler(multivariate=True, group=True)
 
 
+def run_tpe(k: int, sequence_dict: Dict[int, List[int]], hash_dict: Dict[int, int]) -> None:
+    hash_dict[k] = hash("nondeterministic hash")
+    sampler = TPESampler(n_startup_trials=1, seed=2, multivariate=True, group=True)
+    study = create_study(sampler=sampler)
+    sequence = []
+    for _ in range(10):
+        trial = study.ask()
+        picked = [i for i in range(10) if trial.suggest_int(str(i), 0, 1) == 1]
+        study.tell(trial, len(picked))
+        sequence.extend(picked)
+    sequence_dict[k] = sequence
+
+
+# Test deterministic iteration as discussed in https://github.com/optuna/optuna/issues/3137
+def test_group_deterministic_iteration() -> None:
+    from multiprocessing import Manager
+    from multiprocessing import Process
path: tests/samplers_tests/tpe_tests/test_sampler.py
text: @user1:
How about moving these imports to the top of this file?

diff_hunk: @@ -1054,6 +1056,37 @@ def test_group_experimental_warning() -> None:
         _ = TPESampler(multivariate=True, group=True)
 
 
+def run_tpe(k: int, sequence_dict: Dict[int, List[int]], hash_dict: Dict[int, int]) -> None:
+    hash_dict[k] = hash("nondeterministic hash")
+    sampler = TPESampler(n_startup_trials=1, seed=2, multivariate=True, group=True)
+    study = create_study(sampler=sampler)
+    sequence = []
+    for _ in range(10):
+        trial = study.ask()
+        picked = [i for i in range(10) if trial.suggest_int(str(i), 0, 1) == 1]
+        study.tell(trial, len(picked))
+        sequence.extend(picked)
+    sequence_dict[k] = sequence
path: tests/samplers_tests/tpe_tests/test_sampler.py
text: @user1:
How about simplifying the test using a standard `Study.optimize`? For example, I locally confirmed that the following code pass with this PR, but fails with the master.
```suggestion
    sampler = TPESampler(n_startup_trials=1, seed=2, multivariate=True, group=True)
    study = create_study(sampler=sampler)
    study.optimize(lambda t: np.sum(t.suggest_int(f"x{i}", 0, 10) for i in range(10)), n_trials=2)
    sequence_dict[k] = list(study.trials[-1].params.values())
```