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,15 @@ 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`."
+            )
+            self._storage = storage
+        else:
+            self._storage = ot.storages.InMemoryStorage()
+
         self._completed_trials = set()
 
         self._ot_trials = {}
@@ -380,7 +399,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 +420,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,

python/ray/tune/tests/test_searchers.py

@@ -232,22 +232,32 @@ def testNevergradWithRequiredOptimizerKwargs(self):
 
     def testOptuna(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)),
+                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 +368,18 @@ 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
 
+        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}],
@@ -373,6 +389,7 @@ def testOptuna(self):
         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,
@@ -608,9 +625,19 @@ def testNevergrad(self):
         self._restore(searcher)
 
     def testOptuna(self):
+        from optuna.storages import JournalStorage
+        from optuna.storages.journal import JournalFileBackend
+
         from ray.tune.search.optuna import OptunaSearch
 
-        searcher = OptunaSearch(space=self.config, metric=self.metric_name, mode="max")
+        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()
@@ -662,15 +689,20 @@ def tearDownClass(cls):
 
     def testOptuna(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 +717,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__":