[tune] OptunaSearch: check compatibility of search space with evaluated_rewards

python/ray/tune/suggest/optuna.py

@@ -54,7 +54,6 @@ class _OptunaTrialSuggestCaptor:
     `suggest_` callables with a function capturing the returned value,
     which will be saved in the ``captured_values`` dict.
     """
-
     def __init__(self, ot_trial: OptunaTrial) -> None:
         self.ot_trial = ot_trial
         self.captured_values: Dict[str, Any] = {}
@@ -126,6 +125,12 @@ class OptunaSearch(Searcher):
             needing to re-compute the trial. Must be the same length as
             points_to_evaluate.
 
+            ..warning::
+                When using evaluated_rewards, the search space `space` must
+                be provided as a :class:`dict` with parameter names as keys and
+                ``optuna.distributions``. The define-by-run definition is not
+                yet supported.
+
     Tune automatically converts search spaces to Optuna's format:
 
     .. code-block:: python
@@ -151,7 +156,7 @@ class OptunaSearch(Searcher):
         from ray.tune.suggest.optuna import OptunaSearch
         import optuna
 
-        config = {
+        space = {
             "a": optuna.distributions.UniformDistribution(6, 8),
             "b": optuna.distributions.LogUniformDistribution(1e-4, 1e-2),
         }
@@ -178,14 +183,58 @@ def define_search_space(trial: optuna.Trial):
 
         tune.run(trainable, search_alg=optuna_search)
 
+    You can pass configs that will be evaluated first using
+    `points_to_evaluate`:
+
+    .. code-block:: python
+
+        from ray.tune.suggest.optuna import OptunaSearch
+        import optuna
+
+        space = {
+            "a": optuna.distributions.UniformDistribution(6, 8),
+            "b": optuna.distributions.LogUniformDistribution(1e-4, 1e-2),
+        }
+
+        optuna_search = OptunaSearch(
+            space,
+            points_to_evaluate=[{"a": 6.5, "b": 5e-4}, {"a": 7.5, "b": 1e-3}]
+            metric="loss",
+            mode="min")
+
+        tune.run(trainable, search_alg=optuna_search)
+
+    Avoid re-running evaluated trials by passing the rewards together with
+    `points_to_evaluate`:
+
+    .. code-block:: python
+
+        from ray.tune.suggest.optuna import OptunaSearch
+        import optuna
+
+        space = {
+            "a": optuna.distributions.UniformDistribution(6, 8),
+            "b": optuna.distributions.LogUniformDistribution(1e-4, 1e-2),
+        }
+
+        optuna_search = OptunaSearch(
+            space,
+            points_to_evaluate=[{"a": 6.5, "b": 5e-4}, {"a": 7.5, "b": 1e-3}]
+            evaluated_rewards=[0.89, 0.42]
+            metric="loss",
+            mode="min")
+
+        tune.run(trainable, search_alg=optuna_search)
+
     .. versionadded:: 0.8.8
 
     """
-
     def __init__(self,
-                 space: Optional[Union[Dict[str, "OptunaDistribution"], List[
-                     Tuple], Callable[["OptunaTrial"], Optional[Dict[
-                         str, Any]]]]] = None,
+                 space: Optional[Union[Dict[str,
+                                            "OptunaDistribution"], List[Tuple],
+                                       Callable[["OptunaTrial"],
+                                                Optional[Dict[str,
+                                                              Any]]]]] = None,
                  metric: Optional[str] = None,
                  mode: Optional[str] = None,
                  points_to_evaluate: Optional[List[Dict]] = None,
@@ -194,18 +243,17 @@ def __init__(self,
                  evaluated_rewards: Optional[List] = None):
         assert ot is not None, (
             "Optuna must be installed! Run `pip install optuna`.")
-        super(OptunaSearch, self).__init__(
-            metric=metric,
-            mode=mode,
-            max_concurrent=None,
-            use_early_stopped_trials=None)
+        super(OptunaSearch, self).__init__(metric=metric,
+                                           mode=mode,
+                                           max_concurrent=None,
+                                           use_early_stopped_trials=None)
 
         if isinstance(space, dict) and space:
             resolved_vars, domain_vars, grid_vars = parse_spec_vars(space)
             if domain_vars or grid_vars:
                 logger.warning(
-                    UNRESOLVED_SEARCH_SPACE.format(
-                        par="space", cls=type(self).__name__))
+                    UNRESOLVED_SEARCH_SPACE.format(par="space",
+                                                   cls=type(self).__name__))
                 space = self.convert_search_space(space)
             else:
                 # Flatten to support nested dicts
@@ -287,9 +335,11 @@ def set_search_properties(self, metric: Optional[str], mode: Optional[str],
         self._setup_study(mode)
         return True
 
-    def _suggest_from_define_by_run_func(
-            self, func: Callable[["OptunaTrial"], Optional[Dict[str, Any]]],
-            ot_trial: "OptunaTrial") -> Dict:
+    def _suggest_from_define_by_run_func(self,
+                                         func: Callable[["OptunaTrial"],
+                                                        Optional[Dict[str,
+                                                                      Any]]],
+                                         ot_trial: "OptunaTrial") -> Dict:
         captor = _OptunaTrialSuggestCaptor(ot_trial)
         time_start = time.time()
         ret = func(captor)
@@ -321,14 +371,13 @@ def _suggest_from_define_by_run_func(
     def suggest(self, trial_id: str) -> Optional[Dict]:
         if not self._space:
             raise RuntimeError(
-                UNDEFINED_SEARCH_SPACE.format(
-                    cls=self.__class__.__name__, space="space"))
+                UNDEFINED_SEARCH_SPACE.format(cls=self.__class__.__name__,
+                                              space="space"))
         if not self._metric or not self._mode:
             raise RuntimeError(
-                UNDEFINED_METRIC_MODE.format(
-                    cls=self.__class__.__name__,
-                    metric=self._metric,
-                    mode=self._mode))
+                UNDEFINED_METRIC_MODE.format(cls=self.__class__.__name__,
+                                             metric=self._metric,
+                                             mode=self._mode))
 
         if isinstance(self._space, list):
             # Keep for backwards compatibility
@@ -340,8 +389,8 @@ def suggest(self, trial_id: str) -> Optional[Dict]:
 
             # getattr will fetch the trial.suggest_ function on Optuna trials
             params = {
-                args[0] if len(args) > 0 else kwargs["name"]: getattr(
-                    ot_trial, fn)(*args, **kwargs)
+                args[0] if len(args) > 0 else kwargs["name"]:
+                getattr(ot_trial, fn)(*args, **kwargs)
                 for (fn, args, kwargs) in self._space
             }
         elif callable(self._space):
@@ -394,14 +443,18 @@ def add_evaluated_point(self,
                             intermediate_values: Optional[List[float]] = None):
         if not self._space:
             raise RuntimeError(
-                UNDEFINED_SEARCH_SPACE.format(
-                    cls=self.__class__.__name__, space="space"))
+                UNDEFINED_SEARCH_SPACE.format(cls=self.__class__.__name__,
+                                              space="space"))
         if not self._metric or not self._mode:
             raise RuntimeError(
-                UNDEFINED_METRIC_MODE.format(
-                    cls=self.__class__.__name__,
-                    metric=self._metric,
-                    mode=self._mode))
+                UNDEFINED_METRIC_MODE.format(cls=self.__class__.__name__,
+                                             metric=self._metric,
+                                             mode=self._mode))
+        if callable(self._space):
+            raise TypeError(
+                "Define-by-run function passed in `space` argument is not"
+                "yet supported when using `evaluated_rewards`. Please provide"
+                "a 'OptunaDistribution' dict.")
 
         ot_trial_state = OptunaTrialState.COMPLETE
         if error: