Repository: optuna/optuna

Issue:
`Study.optimize` cannot be stopped properly when `sampler=GridSampler` and `pruner=HyperbandPruner`
## Expected behavior

The study should be stopped without errors.

## Environment

- Optuna version: v2.6.0.dev
- Python version: 3.8.5
- OS: Ubuntu 20.04
- (Optional) Other libraries and their versions: N/A

## Error messages, stack traces, or logs

```console
[I 2021-02-08 10:06:01,525] Trial 22 pruned. 
Traceback (most recent call last):
  File "study-error-2322.py", line 18, in <module>
    study.optimize(objective, n_trials=2 ** 4 + 10)
  File "/home/yanase/pfn/code/optuna/optuna/study.py", line 378, in optimize
    _optimize(
  File "/home/yanase/pfn/code/optuna/optuna/_optimize.py", line 66, in _optimize
    _optimize_sequential(
  File "/home/yanase/pfn/code/optuna/optuna/_optimize.py", line 163, in _optimize_sequential
    trial = _run_trial(study, func, catch)
  File "/home/yanase/pfn/code/optuna/optuna/_optimize.py", line 237, in _run_trial
    study.tell(trial, values=values, state=state)
  File "/home/yanase/pfn/code/optuna/optuna/study.py", line 627, in tell
    self.sampler.after_trial(study, frozen_trial, state, values)
  File "/home/yanase/pfn/code/optuna/optuna/samplers/_grid.py", line 196, in after_trial
    study.stop()
  File "/home/yanase/pfn/code/optuna/optuna/pruners/_hyperband.py", line 306, in __getattribute__
    raise AttributeError(
AttributeError: _BracketStudy does not have attribute of 'stop'
```


## Reproducible examples (optional)

```python
import optuna


def objective(trial):
    total = 0
    for i in range(4):
        x = trial.suggest_categorical(f"x{i}", (0, 1))
        total += x
        trial.report(total, i)
        if trial.should_prune():
            raise optuna.TrialPruned()
    return total


if __name__ == "__main__":
    sampler = optuna.samplers.GridSampler(search_space={f"x{i}": (0, 1) for i in range(4)})
    study = optuna.create_study(sampler=sampler, pruner=optuna.pruners.HyperbandPruner())
    study.optimize(objective, n_trials=2 ** 4 + 10)
    print(study.trials_dataframe())
```

## Additional context (optional)

### Cause

This is mainly because `_BracketStudy` does not have the `stop` attribute since it was filtered out by `_VALID_ATTRS`:

https://github.com/optuna/optuna/blob/533f081d33817f2a9a37b21916eb5f03a0fc13c2/optuna/pruners/_hyperband.py#L271-L283

When I simply add `stop` to `_VALID_ATTRS`, another error occured as follows:

```console
[I 2021-02-08 10:07:20,648] Trial 15 finished with value: 4.0 and parameters: {'x0': 1, 'x1': 1, 'x2': 1, 'x3': 1}. Best is trial 2 with value: 0.0.
Traceback (most recent call last):
  File "study-error-2322.py", line 18, in <module>
    study.optimize(objective, n_trials=2 ** 4 + 10)
  File "/home/yanase/pfn/code/optuna/optuna/study.py", line 378, in optimize
    _optimize(
  File "/home/yanase/pfn/code/optuna/optuna/_optimize.py", line 66, in _optimize
    _optimize_sequential(
  File "/home/yanase/pfn/code/optuna/optuna/_optimize.py", line 163, in _optimize_sequential
    trial = _run_trial(study, func, catch)
  File "/home/yanase/pfn/code/optuna/optuna/_optimize.py", line 237, in _run_trial
    study.tell(trial, values=values, state=state)
  File "/home/yanase/pfn/code/optuna/optuna/study.py", line 627, in tell
    self.sampler.after_trial(study, frozen_trial, state, values)
  File "/home/yanase/pfn/code/optuna/optuna/samplers/_grid.py", line 196, in after_trial
    study.stop()
  File "/home/yanase/pfn/code/optuna/optuna/study.py", line 785, in stop
    if self._optimize_lock.acquire(False):
  File "/home/yanase/pfn/code/optuna/optuna/pruners/_hyperband.py", line 307, in __getattribute__
    raise AttributeError(
AttributeError: _BracketStudy does not have attribute of '_optimize_lock'
```

Further investigation is required.

### CI

The error can be seen in the [CI log](https://github.com/optuna/optuna/pull/2322/checks?check_run_id=1849119156). This error occurs stochastically.

```console
>       study.optimize(objective, n_trials=10)

tests/pruners_tests/test_hyperband.py:213: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
optuna/study.py:387: in optimize
    show_progress_bar=show_progress_bar,
optuna/_optimize.py:76: in _optimize
    progress_bar=progress_bar,
optuna/_optimize.py:163: in _optimize_sequential
    trial = _run_trial(study, func, catch)
optuna/_optimize.py:237: in _run_trial
    study.tell(trial, values=values, state=state)
optuna/study.py:627: in tell
    self.sampler.after_trial(study, frozen_trial, state, values)
optuna/samplers/_grid.py:196: in after_trial
    study.stop()
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

self = <optuna.pruners._hyperband.HyperbandPruner._create_bracket_study.<locals>._BracketStudy object at 0x7fa8ced1d6d0>
attr_name = 'stop'

    def __getattribute__(self, attr_name):  # type: ignore
        if attr_name not in _BracketStudy._VALID_ATTRS:
            raise AttributeError(
>               "_BracketStudy does not have attribute of '{}'".format(attr_name)
            )
E           AttributeError: _BracketStudy does not have attribute of 'stop'

optuna/pruners/_hyperband.py:307: AttributeError
```



Discussion:

Resolving patch:
diff --git a/optuna/pruners/_hyperband.py b/optuna/pruners/_hyperband.py
--- a/optuna/pruners/_hyperband.py
+++ b/optuna/pruners/_hyperband.py
@@ -280,6 +280,8 @@ class _BracketStudy(optuna.study.Study):
                 "sampler",
                 "trials",
                 "_is_multi_objective",
+                "stop",
+                "_study",
             )
 
             def __init__(self, study: "optuna.study.Study", bracket_id: int) -> None:
@@ -289,6 +291,7 @@ def __init__(self, study: "optuna.study.Study", bracket_id: int) -> None:
                     sampler=study.sampler,
                     pruner=study.pruner,
                 )
+                self._study = study
                 self._bracket_id = bracket_id
 
             def get_trials(
@@ -301,6 +304,11 @@ def get_trials(
                 assert isinstance(pruner, HyperbandPruner)
                 return [t for t in trials if pruner._get_bracket_id(self, t) == self._bracket_id]
 
+            def stop(self) -> None:
+                # `stop` should stop the original study's optimization loop instead of
+                # `_BracketStudy`.
+                self._study.stop()
+
             def __getattribute__(self, attr_name):  # type: ignore
                 if attr_name not in _BracketStudy._VALID_ATTRS:
                     raise AttributeError(
diff --git a/optuna/samplers/_grid.py b/optuna/samplers/_grid.py
--- a/optuna/samplers/_grid.py
+++ b/optuna/samplers/_grid.py
@@ -211,7 +211,13 @@ def _get_unvisited_grid_ids(self, study: Study) -> List[int]:
 
         # List up unvisited grids based on already finished ones.
         visited_grids = []
-        for t in study.trials:
+
+        # We directly query the storage to get trials here instead of `study.get_trials`,
+        # since some pruners such as `HyperbandPruner` use the study transformed
+        # to filter trials. See https://github.com/optuna/optuna/issues/2327 for details.
+        trials = study._storage.get_all_trials(study._study_id, deepcopy=False)
+
+        for t in trials:
             if (
                 t.state.is_finished()
                 and "grid_id" in t.system_attrs


Test patch:
diff --git a/tests/pruners_tests/test_hyperband.py b/tests/pruners_tests/test_hyperband.py
--- a/tests/pruners_tests/test_hyperband.py
+++ b/tests/pruners_tests/test_hyperband.py
@@ -216,3 +216,27 @@ def objective(trial: optuna.trial.Trial) -> float:
 def test_incompatibility_between_bootstrap_count_and_auto_max_resource() -> None:
     with pytest.raises(ValueError):
         optuna.pruners.HyperbandPruner(max_resource="auto", bootstrap_count=1)
+
+
+def test_hyperband_pruner_and_grid_sampler() -> None:
+    pruner = optuna.pruners.HyperbandPruner(
+        min_resource=MIN_RESOURCE, max_resource=MAX_RESOURCE, reduction_factor=REDUCTION_FACTOR
+    )
+
+    search_space = {"x": [-50, 0, 50], "y": [-99, 0, 99]}
+    sampler = optuna.samplers.GridSampler(search_space)
+
+    study = optuna.study.create_study(sampler=sampler, pruner=pruner)
+
+    def objective(trial: optuna.trial.Trial) -> float:
+        for i in range(N_REPORTS):
+            trial.report(i, step=i)
+
+        x = trial.suggest_float("x", -100, 100)
+        y = trial.suggest_int("y", -100, 100)
+        return x ** 2 + y ** 2
+
+    study.optimize(objective, n_trials=10)
+
+    trials = study.trials
+    assert len(trials) == 9