Repository: facebookresearch/hydra

Issue:
Support optuna.samplers.GridSampler
Hi!

I tried to use `optuna.samplers.GridSampler` for some model with this example:

```yaml
defaults:
  - override /hydra/sweeper: optuna

hydra:
  sweeper:
    _target_: hydra_plugins.hydra_optuna_sweeper.optuna_sweeper.OptunaSweeper
    storage: null
    study_name: null
    n_jobs: 1

    sampler:
      _target_: optuna.samplers.GridSampler
      search_space: {"model.n_d": [ 4, 8, 16 ]}
```
but I got this error:

```
Key 'search_space' not in 'TPESamplerConfig'
    full_key: hydra.sweeper.sampler.search_space
    reference_type=SamplerConfig
    object_type=TPESamplerConfig
```

It is strange, because I create GridSampler, not TPESampler.


## System information
- **Hydra Version** : 
hydra-core==1.1.0
hydra-optuna-sweeper==1.1.0 
- **Python version** : 3.8 
- **Operating system** : Windows 10



Discussion:
As far as I understood, `GridSampler` isn't mentioned here:
https://github.com/facebookresearch/hydra/blob/main/plugins/hydra_optuna_sweeper/hydra_plugins/hydra_optuna_sweeper/config.py

Is this sampler supported?
hi @GillianGrayson 
you are right, `GridSampler`  is not supported. Otherwise it'd be listed in the config.py.

Supporting this should be relatively straightforward (it would require adding a [config dataclass](https://github.com/facebookresearch/hydra/blob/main/plugins/hydra_optuna_sweeper/hydra_plugins/hydra_optuna_sweeper/config.py#L27-L42) and [registering it with the config store](https://github.com/facebookresearch/hydra/blob/main/plugins/hydra_optuna_sweeper/hydra_plugins/hydra_optuna_sweeper/config.py#L199-L204))


Update:
if you are intersted in contributing, pls check [my comment](https://github.com/facebookresearch/hydra/pull/1840#issuecomment-933822775) here  for more info

While `GridSampler` is not supported yet, Hydra's `BasicSweeper` actually supports grid search out of the box with a rich override syntax available - the only thing missing is you cannot define the sweep in a config file (which we will be evaluating for next Hydra release.)

https://hydra.cc/docs/tutorials/basic/running_your_app/multi-run
https://hydra.cc/docs/advanced/override_grammar/basic

cc @GillianGrayson 

Resolving patch:
diff --git a/plugins/hydra_optuna_sweeper/hydra_plugins/hydra_optuna_sweeper/_impl.py b/plugins/hydra_optuna_sweeper/hydra_plugins/hydra_optuna_sweeper/_impl.py
--- a/plugins/hydra_optuna_sweeper/hydra_plugins/hydra_optuna_sweeper/_impl.py
+++ b/plugins/hydra_optuna_sweeper/hydra_plugins/hydra_optuna_sweeper/_impl.py
@@ -1,4 +1,5 @@
 # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+import functools
 import logging
 import sys
 import warnings
@@ -134,6 +135,7 @@ def create_params_from_overrides(
     parsed = parser.parse_overrides(arguments)
     search_space_distributions = dict()
     fixed_params = dict()
+
     for override in parsed:
         param_name = override.get_key_element()
         value = create_optuna_distribution_from_override(override)
@@ -264,6 +266,21 @@ def _parse_sweeper_params_config(self) -> List[str]:
             params_conf.append(f"{k}={v}")
         return params_conf
 
+    def _to_grid_sampler_choices(self, distribution: BaseDistribution) -> Any:
+        if isinstance(distribution, CategoricalDistribution):
+            return distribution.choices
+        elif isinstance(distribution, IntUniformDistribution):
+            assert (
+                distribution.step is not None
+            ), "`step` of IntUniformDistribution must be a positive integer."
+            n_items = (distribution.high - distribution.low) // distribution.step
+            return [distribution.low + i * distribution.step for i in range(n_items)]
+        elif isinstance(distribution, DiscreteUniformDistribution):
+            n_items = int((distribution.high - distribution.low) // distribution.q)
+            return [distribution.low + i * distribution.q for i in range(n_items)]
+        else:
+            raise ValueError("GridSampler only supports discrete distributions.")
+
     def sweep(self, arguments: List[str]) -> None:
         assert self.config is not None
         assert self.launcher is not None
@@ -274,10 +291,31 @@ def sweep(self, arguments: List[str]) -> None:
         self._process_searchspace_config()
         params_conf = self._parse_sweeper_params_config()
         params_conf.extend(arguments)
+
+        is_grid_sampler = (
+            isinstance(self.sampler, functools.partial)
+            and self.sampler.func == optuna.samplers.GridSampler  # type: ignore
+        )
+
         search_space_distributions, fixed_params = create_params_from_overrides(
             params_conf
         )
 
+        if is_grid_sampler:
+            search_space_for_grid_sampler = {
+                name: self._to_grid_sampler_choices(distribution)
+                for name, distribution in search_space_distributions.items()
+            }
+
+            self.sampler = self.sampler(search_space_for_grid_sampler)
+            n_trial = 1
+            for v in search_space_for_grid_sampler.values():
+                n_trial *= len(v)
+            self.n_trials = min(self.n_trials, n_trial)
+            log.info(
+                f"Updating num of trials to {self.n_trials} due to using GridSampler."
+            )
+
         # Remove fixed parameters from Optuna search space.
         for param_name in fixed_params:
             if param_name in search_space_distributions:
@@ -334,7 +372,19 @@ def sweep(self, arguments: List[str]) -> None:
                                 "The number of the values and the number of the objectives are"
                                 f" mismatched. Expect {len(directions)}, but actually {len(values)}."
                             )
-                    study.tell(trial=trial, state=state, values=values)
+
+                    try:
+                        study.tell(trial=trial, state=state, values=values)
+                    except RuntimeError as e:
+                        if (
+                            is_grid_sampler
+                            and "`Study.stop` is supposed to be invoked inside an objective function or a callback."
+                            in str(e)
+                        ):
+                            pass
+                        else:
+                            raise e
+
                 except Exception as e:
                     state = optuna.trial.TrialState.FAIL
                     study.tell(trial=trial, state=state, values=values)
diff --git a/plugins/hydra_optuna_sweeper/hydra_plugins/hydra_optuna_sweeper/config.py b/plugins/hydra_optuna_sweeper/hydra_plugins/hydra_optuna_sweeper/config.py
--- a/plugins/hydra_optuna_sweeper/hydra_plugins/hydra_optuna_sweeper/config.py
+++ b/plugins/hydra_optuna_sweeper/hydra_plugins/hydra_optuna_sweeper/config.py
@@ -21,7 +21,17 @@ class Direction(Enum):
 @dataclass
 class SamplerConfig:
     _target_: str = MISSING
-    seed: Optional[int] = None
+
+
+@dataclass
+class GridSamplerConfig(SamplerConfig):
+    """
+    https://optuna.readthedocs.io/en/stable/reference/generated/optuna.samplers.GridSampler.html
+    """
+
+    _target_: str = "optuna.samplers.GridSampler"
+    # search_space will be populated at run time based on hydra.sweeper.params
+    _partial_: bool = True
 
 
 @dataclass
@@ -31,6 +41,7 @@ class TPESamplerConfig(SamplerConfig):
     """
 
     _target_: str = "optuna.samplers.TPESampler"
+    seed: Optional[int] = None
 
     consider_prior: bool = True
     prior_weight: float = 1.0
@@ -49,6 +60,7 @@ class RandomSamplerConfig(SamplerConfig):
     """
 
     _target_: str = "optuna.samplers.RandomSampler"
+    seed: Optional[int] = None
 
 
 @dataclass
@@ -58,6 +70,7 @@ class CmaEsSamplerConfig(SamplerConfig):
     """
 
     _target_: str = "optuna.samplers.CmaEsSampler"
+    seed: Optional[int] = None
 
     x0: Optional[Dict[str, Any]] = None
     sigma0: Optional[float] = None
@@ -77,6 +90,7 @@ class NSGAIISamplerConfig(SamplerConfig):
     """
 
     _target_: str = "optuna.samplers.NSGAIISampler"
+    seed: Optional[int] = None
 
     population_size: int = 50
     mutation_prob: Optional[float] = None
@@ -92,6 +106,7 @@ class MOTPESamplerConfig(SamplerConfig):
     """
 
     _target_: str = "optuna.samplers.MOTPESampler"
+    seed: Optional[int] = None
 
     consider_prior: bool = True
     prior_weight: float = 1.0
@@ -210,3 +225,10 @@ class OptunaSweeperConf:
     node=MOTPESamplerConfig,
     provider="optuna_sweeper",
 )
+
+ConfigStore.instance().store(
+    group="hydra/sweeper/sampler",
+    name="grid",
+    node=GridSamplerConfig,
+    provider="optuna_sweeper",
+)


Test patch:
diff --git a/plugins/hydra_optuna_sweeper/tests/conf/test_grid.yaml b/plugins/hydra_optuna_sweeper/tests/conf/test_grid.yaml
new file mode 100644
--- /dev/null
+++ b/plugins/hydra_optuna_sweeper/tests/conf/test_grid.yaml
@@ -0,0 +1,19 @@
+defaults:
+  - override hydra/sweeper: optuna
+  - override hydra/sweeper/sampler: grid
+
+hydra:
+  sweeper:
+    direction: minimize
+    study_name: sphere
+    storage: null
+    n_trials: 20
+    n_jobs: 1
+    params:
+      x: choice(-1, 1)
+      y: range(-1.0, 1.0, step=1)
+      z: choice("foo", "bar")
+
+x: 1
+y: 1
+z: foo
diff --git a/plugins/hydra_optuna_sweeper/tests/test_optuna_sweeper_plugin.py b/plugins/hydra_optuna_sweeper/tests/test_optuna_sweeper_plugin.py
--- a/plugins/hydra_optuna_sweeper/tests/test_optuna_sweeper_plugin.py
+++ b/plugins/hydra_optuna_sweeper/tests/test_optuna_sweeper_plugin.py
@@ -192,6 +192,41 @@ def test_optuna_example(with_commandline: bool, tmpdir: Path) -> None:
     assert returns["best_value"] <= 2.27
 
 
+@mark.parametrize("num_trials", (10, 1))
+def test_example_with_grid_sampler(
+    tmpdir: Path,
+    num_trials: int,
+) -> None:
+    storage = "sqlite:///" + os.path.join(str(tmpdir), "test.db")
+    study_name = "test-grid-sampler"
+    cmd = [
+        "example/sphere.py",
+        "--multirun",
+        "--config-dir=tests/conf",
+        "--config-name=test_grid",
+        "hydra.sweep.dir=" + str(tmpdir),
+        "hydra.job.chdir=False",
+        f"hydra.sweeper.n_trials={num_trials}",
+        "hydra.sweeper.n_jobs=1",
+        f"hydra.sweeper.storage={storage}",
+        f"hydra.sweeper.study_name={study_name}",
+    ]
+    run_python_script(cmd)
+    returns = OmegaConf.load(f"{tmpdir}/optimization_results.yaml")
+    assert isinstance(returns, DictConfig)
+    bv, bx, by, bz = (
+        returns["best_value"],
+        returns["best_params"]["x"],
+        returns["best_params"]["y"],
+        returns["best_params"]["z"],
+    )
+    if num_trials >= 12:
+        assert bv == 1 and abs(bx) == 1 and by == 0
+    else:
+        assert bx in [-1, 1] and by in [-1, 0]
+    assert bz in ["foo", "bar"]
+
+
 @mark.parametrize("with_commandline", (True, False))
 def test_optuna_multi_objective_example(with_commandline: bool, tmpdir: Path) -> None:
     cmd = [