diff --git a/python/ray/tune/search/optuna/optuna_search.py b/python/ray/tune/search/optuna/optuna_search.py index f1656039a4af..3cc1972483ca 100644 --- a/python/ray/tune/search/optuna/optuna_search.py +++ b/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, diff --git a/python/ray/tune/tests/test_searchers.py b/python/ray/tune/tests/test_searchers.py index c7b88f87109b..5346966c81d4 100644 --- a/python/ray/tune/tests/test_searchers.py +++ b/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__": diff --git a/python/requirements/ml/tune-requirements.txt b/python/requirements/ml/tune-requirements.txt index 03ebfbd876b6..cefcb903f962 100644 --- a/python/requirements/ml/tune-requirements.txt +++ b/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