From 882f7d3863059dfb361f44652815179d1dc5d1ec Mon Sep 17 00:00:00 2001
From: Carl Assmann <assmann@hey.com>
Date: Thu, 23 Sep 2021 17:00:11 +0200
Subject: [PATCH] [tune] OptunaSearch: check compatibility of search space with
 evaluated_rewards (#18625)

Co-authored-by: Antoni Baum <antoni.baum@protonmail.com>
Co-authored-by: Kai Fricke <krfricke@users.noreply.github.com>
---
 python/ray/tune/suggest/optuna.py       | 58 ++++++++++++++++++++++++-
 python/ray/tune/tests/test_searchers.py | 13 ++++++
 2 files changed, 70 insertions(+), 1 deletion(-)

diff --git a/python/ray/tune/suggest/optuna.py b/python/ray/tune/suggest/optuna.py
index 43ef5b2e4..c69ef4fa3 100644
--- a/python/ray/tune/suggest/optuna.py
+++ b/python/ray/tune/suggest/optuna.py
@@ -114,6 +114,13 @@ 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`` instances as values. The
+                define-by-run search space definition is not yet supported with
+                this functionality.
+
     Tune automatically converts search spaces to Optuna's format:
 
     .. code-block:: python
@@ -139,7 +146,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),
         }
@@ -166,6 +173,49 @@ class OptunaSearch(Searcher):
 
         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
 
     """
@@ -367,6 +417,12 @@ class OptunaSearch(Searcher):
                     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 "
+                "an `OptunaDistribution` dict or pass a Ray Tune "
+                "search space to `tune.run()`.")
 
         ot_trial_state = OptunaTrialState.COMPLETE
         if error:
diff --git a/python/ray/tune/tests/test_searchers.py b/python/ray/tune/tests/test_searchers.py
index cf46a1d76..a428b6f16 100644
--- a/python/ray/tune/tests/test_searchers.py
+++ b/python/ray/tune/tests/test_searchers.py
@@ -317,6 +317,19 @@ class AddEvaluatedPointTest(unittest.TestCase):
         self.assertTrue(
             searcher._ot_study.trials[-1].state == TrialState.PRUNED)
 
+        def dbr_space(trial):
+            return {
+                self.param_name: trial.suggest_float(self.param_name, 0.0, 5.0)
+            }
+
+        dbr_searcher = OptunaSearch(
+            space=dbr_space,
+            metric="metric",
+            mode="max",
+        )
+        with self.assertRaises(TypeError):
+            dbr_searcher.add_evaluated_point(point, 1.0)
+
     def testHEBO(self):
         from ray.tune.suggest.hebo import HEBOSearch