updated optuna_search to allow users to configure optuna storage

python/ray/tune/search/optuna/optuna_search.py

@@ -31,12 +31,14 @@
     import optuna as ot
     from optuna.distributions import BaseDistribution as OptunaDistribution
     from optuna.samplers import BaseSampler
+    from optuna.storages import BaseStorage
     from optuna.trial import Trial as OptunaTrial
     from optuna.trial import TrialState as OptunaTrialState
 except ImportError:
     ot = None
     OptunaDistribution = None
     BaseSampler = None
+    BaseStorage = None
     OptunaTrialState = None
     OptunaTrial = None
 
@@ -133,7 +135,11 @@ class OptunaSearch(Searcher):
                 a delay when suggesting new configurations.
                 This is an Optuna issue and may be fixed in a future
                 Optuna release.
-
+        study_name: Optuna study name that uniquely identifies the trial
+            results. Defaults to ``"optuna"``.
+        storage: Optuna storage used for storing trial results to
+            storages other than in-memory storage,
+            for instance optuna.storages.RDBStorage.
         seed: Seed to initialize sampler with. This parameter is only
             used when ``sampler=None``. In all other cases, the sampler
             you pass should be initialized with the seed already.
@@ -322,6 +328,8 @@ def __init__(
         mode: Optional[Union[str, List[str]]] = None,
         points_to_evaluate: Optional[List[Dict]] = None,
         sampler: Optional["BaseSampler"] = None,
+        study_name: Optional[str] = None,
+        storage: Optional["BaseStorage"] = None,
         seed: Optional[int] = None,
         evaluated_rewards: Optional[List] = None,
     ):
@@ -343,8 +351,10 @@ def __init__(
 
         self._points_to_evaluate = points_to_evaluate or []
         self._evaluated_rewards = evaluated_rewards
-
-        self._study_name = "optuna"  # Fixed study name for in-memory storage
+        if study_name:
+            self._study_name = study_name
+        else:
+            self._study_name = "optuna"  # Fixed study name for in-memory storage
 
         if sampler and seed:
             logger.warning(
@@ -362,6 +372,17 @@ def __init__(
         self._sampler = sampler
         self._seed = seed
 
+        if storage:
+            assert isinstance(storage, BaseStorage), (
+                "The `storage` parameter in `OptunaSearcher` must be an instance "
+                "of `optuna.samplers.BaseStorage`."
+            )
+        # If storage is not provided, just set self._storage to None
+        # so that the default in-memory storage is used. We don't explicitly
+        # set self._storage to InMemoryStorage because it is a new API in Optuna
+        # and we want to keep the backward compatibility.
+        self._storage = storage
+
         self._completed_trials = set()
 
         self._ot_trials = {}
@@ -380,7 +401,6 @@ def _setup_study(self, mode: Union[str, list]):
             self._metric = DEFAULT_METRIC
 
         pruner = ot.pruners.NopPruner()
-        storage = ot.storages.InMemoryStorage()
 
         if self._sampler:
             sampler = self._sampler
@@ -402,7 +422,7 @@ def _setup_study(self, mode: Union[str, list]):
             )
 
         self._ot_study = ot.study.create_study(
-            storage=storage,
+            storage=self._storage,
             sampler=sampler,
             pruner=pruner,
             study_name=self._study_name,
@@ -595,6 +615,15 @@ def add_evaluated_point(
         else:
             intermediate_values_dict = None
 
+        # If the trial state is FAILED, the value must be `None` in Optuna==4.1.0
+        # Reference: https://github.com/optuna/optuna/pull/5211
+        # This is a temporary fix for the issue that Optuna enforces the value
+        # to be `None` if the trial state is FAILED.
+        # TODO (hpguo): A better solution may requires us to update the base class
+        # to allow the `value` arg in `add_evaluated_point` being `Optional[float]`.
+        if ot_trial_state == OptunaTrialState.FAIL:
+            value = None
+
         trial = ot.trial.create_trial(
             state=ot_trial_state,
             value=value,

python/ray/tune/tests/test_searchers.py

@@ -240,7 +240,7 @@ def testOptuna(self):
         with self.check_searcher_checkpoint_errors_scope():
             out = tune.run(
                 _invalid_objective,
-                search_alg=OptunaSearch(sampler=RandomSampler(seed=1234)),
+                search_alg=OptunaSearch(sampler=RandomSampler(seed=1234), storage=None),
                 config=self.config,
                 metric="_metric",
                 mode="max",
@@ -249,6 +249,35 @@ def testOptuna(self):
             )
         self.assertCorrectExperimentOutput(out)
 
+    def testOptunaWithStorage(self):
+        from optuna.samplers import RandomSampler
+        from optuna.storages import JournalStorage
+        from optuna.storages.journal import JournalFileBackend
+
+        from ray.tune.search.optuna import OptunaSearch
+
+        np.random.seed(1000)  # At least one nan, inf, -inf and float
+        storage_file_path = "/tmp/my_test_study.log"
+
+        with self.check_searcher_checkpoint_errors_scope():
+            out = tune.run(
+                _invalid_objective,
+                search_alg=OptunaSearch(
+                    sampler=RandomSampler(seed=1234),
+                    study_name="my_test_study",
+                    storage=JournalStorage(
+                        JournalFileBackend(file_path=storage_file_path)
+                    ),
+                ),
+                config=self.config,
+                metric="_metric",
+                mode="max",
+                num_samples=8,
+                reuse_actors=False,
+            )
+        self.assertCorrectExperimentOutput(out)
+        self.assertTrue(os.path.exists(storage_file_path))
+
     def testOptunaReportTooOften(self):
         from optuna.samplers import RandomSampler
 
@@ -358,12 +387,16 @@ def run_add_evaluated_trials(self, searcher, get_len_X, get_len_y):
         searcher_copy.suggest("1")
 
     def testOptuna(self):
+        from optuna.storages import JournalStorage
+        from optuna.storages.journal import JournalFileBackend
         from optuna.trial import TrialState
 
         from ray.tune.search.optuna import OptunaSearch
 
+        # OptunaSearch with in-memory storage
         searcher = OptunaSearch(
             space=self.space,
+            storage=None,
             metric="metric",
             mode="max",
             points_to_evaluate=[{self.param_name: self.valid_value}],
@@ -374,6 +407,23 @@ def testOptuna(self):
 
         self.assertGreater(get_len(searcher), 0)
 
+        # OptunaSearch with external storage
+        storage_file_path = "/tmp/my_test_study.log"
+        searcher = OptunaSearch(
+            space=self.space,
+            study_name="my_test_study",
+            storage=JournalStorage(JournalFileBackend(file_path=storage_file_path)),
+            metric="metric",
+            mode="max",
+            points_to_evaluate=[{self.param_name: self.valid_value}],
+            evaluated_rewards=[1.0],
+        )
+
+        get_len = lambda s: len(s._ot_study.trials)  # noqa E731
+
+        self.assertGreater(get_len(searcher), 0)
+        self.assertTrue(os.path.exists(storage_file_path))
+
         searcher = OptunaSearch(
             space=self.space,
             metric="metric",
@@ -610,13 +660,40 @@ def testNevergrad(self):
     def testOptuna(self):
         from ray.tune.search.optuna import OptunaSearch
 
-        searcher = OptunaSearch(space=self.config, metric=self.metric_name, mode="max")
+        searcher = OptunaSearch(
+            space=self.config,
+            storage=None,
+            metric=self.metric_name,
+            mode="max",
+        )
+        self._save(searcher)
+
+        searcher = OptunaSearch()
+        self._restore(searcher)
+
+        assert "not_completed" in searcher._ot_trials
+
+    def testOptunaWithStorage(self):
+        from optuna.storages import JournalStorage
+        from optuna.storages.journal import JournalFileBackend
+
+        from ray.tune.search.optuna import OptunaSearch
+
+        storage_file_path = "/tmp/my_test_study.log"
+        searcher = OptunaSearch(
+            space=self.config,
+            study_name="my_test_study",
+            storage=JournalStorage(JournalFileBackend(file_path=storage_file_path)),
+            metric=self.metric_name,
+            mode="max",
+        )
         self._save(searcher)
 
         searcher = OptunaSearch()
         self._restore(searcher)
 
         assert "not_completed" in searcher._ot_trials
+        self.assertTrue(os.path.exists(storage_file_path))
 
     def testZOOpt(self):
         from ray.tune.search.zoopt import ZOOptSearch
@@ -671,6 +748,38 @@ def testOptuna(self):
             _multi_objective,
             search_alg=OptunaSearch(
                 sampler=RandomSampler(seed=1234),
+                storage=None,
+                metric=["a", "b", "c"],
+                mode=["max", "min", "max"],
+            ),
+            config=self.config,
+            num_samples=16,
+            reuse_actors=False,
+        )
+
+        best_trial_a = out.get_best_trial("a", "max")
+        self.assertGreaterEqual(best_trial_a.config["a"], 0.8)
+        best_trial_b = out.get_best_trial("b", "min")
+        self.assertGreaterEqual(best_trial_b.config["b"], 0.8)
+        best_trial_c = out.get_best_trial("c", "max")
+        self.assertGreaterEqual(best_trial_c.config["c"], 0.8)
+
+    def testOptunaWithStorage(self):
+        from optuna.samplers import RandomSampler
+        from optuna.storages import JournalStorage
+        from optuna.storages.journal import JournalFileBackend
+
+        from ray.tune.search.optuna import OptunaSearch
+
+        np.random.seed(1000)
+        storage_file_path = "/tmp/my_test_study.log"
+
+        out = tune.run(
+            _multi_objective,
+            search_alg=OptunaSearch(
+                sampler=RandomSampler(seed=1234),
+                study_name="my_test_study",
+                storage=JournalStorage(JournalFileBackend(file_path=storage_file_path)),
                 metric=["a", "b", "c"],
                 mode=["max", "min", "max"],
             ),
@@ -685,6 +794,7 @@ def testOptuna(self):
         self.assertGreaterEqual(best_trial_b.config["b"], 0.8)
         best_trial_c = out.get_best_trial("c", "max")
         self.assertGreaterEqual(best_trial_c.config["c"], 0.8)
+        self.assertTrue(os.path.exists(storage_file_path))
 
 
 if __name__ == "__main__":

python/requirements/ml/tune-requirements.txt

@@ -10,4 +10,4 @@ hpbandster==0.7.4; python_version < "3.12"
 hyperopt @ git+https://github.com/hyperopt/hyperopt.git@2504ee61419737e814e2dec2961b15d12775529c
 future
 nevergrad==0.4.3.post7
-optuna==3.2.0
+optuna==4.1.0

python/requirements_compiled.txt

@@ -295,8 +295,6 @@ cloudpickle==2.2.0
     #   tensorflow-probability
 cma==3.2.2
     # via nevergrad
-cmaes==0.10.0
-    # via optuna
 cmdstanpy==1.2.0
     # via prophet
 colorama==0.4.6
@@ -1166,7 +1164,6 @@ numpy==1.26.4
     #   bayesian-optimization
     #   bokeh
     #   cma
-    #   cmaes
     #   cmdstanpy
     #   configspace
     #   contourpy
@@ -1309,7 +1306,7 @@ opt-einsum==3.3.0
     # via
     #   pyro-ppl
     #   tensorflow
-optuna==3.2.0
+optuna==4.1.0
     # via -r /ray/ci/../python/requirements/ml/tune-requirements.txt
 orjson==3.9.10
     # via gradio