from __future__ import annotations import copy from datetime import datetime import math import pickle import random import sys from time import sleep from typing import Any import numpy as np import pytest import optuna from optuna._typing import JSONSerializable from optuna.distributions import CategoricalDistribution from optuna.distributions import FloatDistribution from optuna.exceptions import UpdateFinishedTrialError from optuna.storages import BaseStorage from optuna.storages._base import DEFAULT_STUDY_NAME_PREFIX from optuna.study._frozen import FrozenStudy from optuna.study._study_direction import StudyDirection from optuna.trial import FrozenTrial from optuna.trial import TrialState @pytest.fixture def storage() -> BaseStorage: raise NotImplementedError class StorageTestCase: def test_create_new_study(self, storage: BaseStorage) -> None: study_id = storage.create_new_study(directions=[StudyDirection.MINIMIZE]) frozen_studies = storage.get_all_studies() assert len(frozen_studies) == 1 assert frozen_studies[0]._study_id == study_id assert frozen_studies[0].study_name.startswith(DEFAULT_STUDY_NAME_PREFIX) study_id2 = storage.create_new_study(directions=[StudyDirection.MINIMIZE]) # Study id must be unique. assert study_id != study_id2 frozen_studies = storage.get_all_studies() assert len(frozen_studies) == 2 assert {s._study_id for s in frozen_studies} == {study_id, study_id2} assert all(s.study_name.startswith(DEFAULT_STUDY_NAME_PREFIX) for s in frozen_studies) def test_create_new_study_unique_id(self, storage: BaseStorage) -> None: study_id = storage.create_new_study(directions=[StudyDirection.MINIMIZE]) study_id2 = storage.create_new_study(directions=[StudyDirection.MINIMIZE]) storage.delete_study(study_id2) study_id3 = storage.create_new_study(directions=[StudyDirection.MINIMIZE]) frozen_studies = storage.get_all_studies() assert {s._study_id for s in frozen_studies} == {study_id, study_id3} def test_create_new_study_with_name(self, storage: BaseStorage) -> None: # Generate unique study_name from the current function name. function_name = self.test_create_new_study_with_name.__name__ study_name = function_name study_id = storage.create_new_study( directions=[StudyDirection.MINIMIZE], study_name=study_name ) assert study_name == storage.get_study_name_from_id(study_id) with pytest.raises(optuna.exceptions.DuplicatedStudyError): storage.create_new_study(directions=[StudyDirection.MINIMIZE], study_name=study_name) def test_delete_study(self, storage: BaseStorage) -> None: study_id = storage.create_new_study(directions=[StudyDirection.MINIMIZE]) storage.create_new_trial(study_id) trials = storage.get_all_trials(study_id) assert len(trials) == 1 with pytest.raises(KeyError): # Deletion of non-existent study. storage.delete_study(study_id + 1) storage.delete_study(study_id) study_id = storage.create_new_study(directions=[StudyDirection.MINIMIZE]) trials = storage.get_all_trials(study_id) assert len(trials) == 0 storage.delete_study(study_id) with pytest.raises(KeyError): # Double free. storage.delete_study(study_id) def test_delete_study_after_create_multiple_studies(self, storage: BaseStorage) -> None: study_id1 = storage.create_new_study(directions=[StudyDirection.MINIMIZE]) study_id2 = storage.create_new_study(directions=[StudyDirection.MINIMIZE]) study_id3 = storage.create_new_study(directions=[StudyDirection.MINIMIZE]) storage.delete_study(study_id2) studies = {s._study_id: s for s in storage.get_all_studies()} assert study_id1 in studies assert study_id2 not in studies assert study_id3 in studies def test_get_study_id_from_name_and_get_study_name_from_id(self, storage: BaseStorage) -> None: # Generate unique study_name from the current function name. function_name = self.test_get_study_id_from_name_and_get_study_name_from_id.__name__ study_name = function_name study_id = storage.create_new_study( directions=[StudyDirection.MINIMIZE], study_name=study_name ) # Test existing study. assert storage.get_study_name_from_id(study_id) == study_name assert storage.get_study_id_from_name(study_name) == study_id # Test not existing study. with pytest.raises(KeyError): storage.get_study_id_from_name("dummy-name") with pytest.raises(KeyError): storage.get_study_name_from_id(study_id + 1) def test_set_and_get_study_directions(self, storage: BaseStorage) -> None: for target in [ (StudyDirection.MINIMIZE,), (StudyDirection.MAXIMIZE,), (StudyDirection.MINIMIZE, StudyDirection.MAXIMIZE), (StudyDirection.MAXIMIZE, StudyDirection.MINIMIZE), [StudyDirection.MINIMIZE, StudyDirection.MAXIMIZE], [StudyDirection.MAXIMIZE, StudyDirection.MINIMIZE], ]: study_id = storage.create_new_study(directions=target) def check_get() -> None: got_directions = storage.get_study_directions(study_id) assert got_directions == list(target), ( "Direction of a study should be a tuple of `StudyDirection` objects." ) # Test setting value. check_get() # Test non-existent study. non_existent_study_id = study_id + 1 with pytest.raises(KeyError): storage.get_study_directions(non_existent_study_id) def test_set_and_get_study_user_attrs(self, storage: BaseStorage) -> None: study_id = storage.create_new_study(directions=[StudyDirection.MINIMIZE]) def check_set_and_get(key: str, value: Any) -> None: storage.set_study_user_attr(study_id, key, value) assert storage.get_study_user_attrs(study_id)[key] == value # Test setting value. for key, value in EXAMPLE_ATTRS.items(): check_set_and_get(key, value) assert storage.get_study_user_attrs(study_id) == EXAMPLE_ATTRS # Test overwriting value. check_set_and_get("dataset", "ImageNet") # Non-existent study id. non_existent_study_id = study_id + 1 with pytest.raises(KeyError): storage.get_study_user_attrs(non_existent_study_id) # Non-existent study id. with pytest.raises(KeyError): storage.set_study_user_attr(non_existent_study_id, "key", "value") def test_set_and_get_study_user_attrs_for_floats(self, storage: BaseStorage) -> None: _test_set_and_get_study_user_attrs_for_floats(storage) def test_set_and_get_study_system_attrs(self, storage: BaseStorage) -> None: study_id = storage.create_new_study(directions=[StudyDirection.MINIMIZE]) def check_set_and_get(key: str, value: Any) -> None: storage.set_study_system_attr(study_id, key, value) assert storage.get_study_system_attrs(study_id)[key] == value # Test setting value. for key, value in EXAMPLE_ATTRS.items(): check_set_and_get(key, value) assert storage.get_study_system_attrs(study_id) == EXAMPLE_ATTRS # Test overwriting value. check_set_and_get("dataset", "ImageNet") # Non-existent study id. non_existent_study_id = study_id + 1 with pytest.raises(KeyError): storage.get_study_system_attrs(non_existent_study_id) # Non-existent study id. with pytest.raises(KeyError): storage.set_study_system_attr(non_existent_study_id, "key", "value") def test_set_and_get_study_system_attrs_for_floats(self, storage: BaseStorage) -> None: _test_set_and_get_study_system_attrs_for_floats(storage) def test_study_user_and_system_attrs_confusion(self, storage: BaseStorage) -> None: study_id = storage.create_new_study(directions=[StudyDirection.MINIMIZE]) for key, value in EXAMPLE_ATTRS.items(): storage.set_study_system_attr(study_id, key, value) assert storage.get_study_system_attrs(study_id) == EXAMPLE_ATTRS assert storage.get_study_user_attrs(study_id) == {} study_id = storage.create_new_study(directions=[StudyDirection.MINIMIZE]) for key, value in EXAMPLE_ATTRS.items(): storage.set_study_user_attr(study_id, key, value) assert storage.get_study_user_attrs(study_id) == EXAMPLE_ATTRS assert storage.get_study_system_attrs(study_id) == {} def test_create_new_trial(self, storage: BaseStorage) -> None: def _check_trials( trials: list[FrozenTrial], idx: int, trial_id: int, time_before_creation: datetime, time_after_creation: datetime, ) -> None: assert len(trials) == idx + 1 assert len({t._trial_id for t in trials}) == idx + 1 assert trial_id in {t._trial_id for t in trials} assert {t.number for t in trials} == set(range(idx + 1)) assert all(t.state == TrialState.RUNNING for t in trials) assert all(t.params == {} for t in trials) assert all(t.intermediate_values == {} for t in trials) assert all(t.user_attrs == {} for t in trials) assert all(t.system_attrs == {} for t in trials) assert all( t.datetime_start < time_before_creation for t in trials if t._trial_id != trial_id and t.datetime_start is not None ) assert all( time_before_creation < t.datetime_start < time_after_creation for t in trials if t._trial_id == trial_id and t.datetime_start is not None ) assert all(t.datetime_complete is None for t in trials) assert all(t.value is None for t in trials) study_id = storage.create_new_study(directions=[StudyDirection.MINIMIZE]) n_trial_in_study = 3 for i in range(n_trial_in_study): time_before_creation = datetime.now() sleep(0.001) # Sleep 1ms to avoid faulty assertion on Windows OS. trial_id = storage.create_new_trial(study_id) sleep(0.001) time_after_creation = datetime.now() trials = storage.get_all_trials(study_id) _check_trials(trials, i, trial_id, time_before_creation, time_after_creation) # Create trial in non-existent study. with pytest.raises(KeyError): storage.create_new_trial(study_id + 1) study_id2 = storage.create_new_study(directions=[StudyDirection.MINIMIZE]) for i in range(n_trial_in_study): storage.create_new_trial(study_id2) trials = storage.get_all_trials(study_id2) # Check that the offset of trial.number is zero. assert {t.number for t in trials} == set(range(i + 1)) trials = storage.get_all_trials(study_id) + storage.get_all_trials(study_id2) # Check trial_ids are unique across studies. assert len({t._trial_id for t in trials}) == 2 * n_trial_in_study @pytest.mark.parametrize( "start_time,complete_time", [(datetime.now(), datetime.now()), (datetime(2022, 9, 1), datetime(2022, 9, 2))], ) def test_create_new_trial_with_template_trial( self, storage: BaseStorage, start_time: datetime, complete_time: datetime ) -> None: template_trial = FrozenTrial( state=TrialState.COMPLETE, value=10000, datetime_start=start_time, datetime_complete=complete_time, params={"x": 0.5}, distributions={"x": FloatDistribution(0, 1)}, user_attrs={"foo": "bar"}, system_attrs={"baz": 123}, intermediate_values={1: 10, 2: 100, 3: 1000}, number=55, # This entry is ignored. trial_id=-1, # dummy value (unused). ) def _check_trials(trials: list[FrozenTrial], idx: int, trial_id: int) -> None: assert len(trials) == idx + 1 assert len({t._trial_id for t in trials}) == idx + 1 assert trial_id in {t._trial_id for t in trials} assert {t.number for t in trials} == set(range(idx + 1)) assert all(t.state == template_trial.state for t in trials) assert all(t.params == template_trial.params for t in trials) assert all(t.distributions == template_trial.distributions for t in trials) assert all(t.intermediate_values == template_trial.intermediate_values for t in trials) assert all(t.user_attrs == template_trial.user_attrs for t in trials) assert all(t.system_attrs == template_trial.system_attrs for t in trials) assert all(t.datetime_start == template_trial.datetime_start for t in trials) assert all(t.datetime_complete == template_trial.datetime_complete for t in trials) assert all(t.value == template_trial.value for t in trials) study_id = storage.create_new_study(directions=[StudyDirection.MINIMIZE]) n_trial_in_study = 3 for i in range(n_trial_in_study): trial_id = storage.create_new_trial(study_id, template_trial=template_trial) trials = storage.get_all_trials(study_id) _check_trials(trials, i, trial_id) # Create trial in non-existent study. with pytest.raises(KeyError): storage.create_new_trial(study_id + 1) study_id2 = storage.create_new_study(directions=[StudyDirection.MINIMIZE]) for i in range(n_trial_in_study): storage.create_new_trial(study_id2, template_trial=template_trial) trials = storage.get_all_trials(study_id2) assert {t.number for t in trials} == set(range(i + 1)) trials = storage.get_all_trials(study_id) + storage.get_all_trials(study_id2) # Check trial_ids are unique across studies. assert len({t._trial_id for t in trials}) == 2 * n_trial_in_study def test_get_trial_number_from_id(self, storage: BaseStorage) -> None: # Check if trial_number starts from 0. study_id = storage.create_new_study(directions=[StudyDirection.MINIMIZE]) trial_id = storage.create_new_trial(study_id) assert storage.get_trial_number_from_id(trial_id) == 0 trial_id = storage.create_new_trial(study_id) assert storage.get_trial_number_from_id(trial_id) == 1 with pytest.raises(KeyError): storage.get_trial_number_from_id(trial_id + 1) def test_set_trial_state_values_for_state(self, storage: BaseStorage) -> None: study_id = storage.create_new_study(directions=[StudyDirection.MINIMIZE]) trial_ids = [storage.create_new_trial(study_id) for _ in ALL_STATES] for trial_id, state in zip(trial_ids, ALL_STATES): if state == TrialState.WAITING: continue assert storage.get_trial(trial_id).state == TrialState.RUNNING datetime_start_prev = storage.get_trial(trial_id).datetime_start storage.set_trial_state_values( trial_id, state=state, values=(0.0,) if state.is_finished() else None ) assert storage.get_trial(trial_id).state == state # Repeated state changes to RUNNING should not trigger further datetime_start changes. if state == TrialState.RUNNING: assert storage.get_trial(trial_id).datetime_start == datetime_start_prev if state.is_finished(): assert storage.get_trial(trial_id).datetime_complete is not None else: assert storage.get_trial(trial_id).datetime_complete is None # Non-existent study. with pytest.raises(KeyError): non_existent_trial_id = max(trial_ids) + 1 storage.set_trial_state_values( non_existent_trial_id, state=TrialState.COMPLETE, ) for state in ALL_STATES: if not state.is_finished(): continue trial_id = storage.create_new_trial(study_id) storage.set_trial_state_values(trial_id, state=state, values=(0.0,)) for state2 in ALL_STATES: # Cannot update states of finished trials. values: list[float] | None = None if state2 == TrialState.COMPLETE: values = [0.0] with pytest.raises(UpdateFinishedTrialError): storage.set_trial_state_values(trial_id, state=state2, values=values) def test_set_trial_state_values_for_floats(self, storage: BaseStorage) -> None: _test_set_trial_state_values_for_floats(storage) def test_get_trial_param_and_get_trial_params(self, storage: BaseStorage) -> None: _, study_to_trials = _setup_studies(storage, n_study=2, n_trial=5, seed=1) for _, trial_id_to_trial in study_to_trials.items(): for trial_id, expected_trial in trial_id_to_trial.items(): assert storage.get_trial_params(trial_id) == expected_trial.params for key in expected_trial.params.keys(): assert storage.get_trial_param(trial_id, key) == expected_trial.distributions[ key ].to_internal_repr(expected_trial.params[key]) non_existent_trial_id = ( max(tid for ts in study_to_trials.values() for tid in ts.keys()) + 1 ) with pytest.raises(KeyError): storage.get_trial_params(non_existent_trial_id) with pytest.raises(KeyError): storage.get_trial_param(non_existent_trial_id, "paramA") existent_trial_id = non_existent_trial_id - 1 with pytest.raises(KeyError): storage.get_trial_param(existent_trial_id, "dummy-key") def test_set_trial_param(self, storage: BaseStorage) -> None: # Setup test across multiple studies and trials. study_id = storage.create_new_study(directions=[StudyDirection.MINIMIZE]) trial_id_1 = storage.create_new_trial(study_id) trial_id_2 = storage.create_new_trial(study_id) trial_id_3 = storage.create_new_trial( storage.create_new_study(directions=[StudyDirection.MINIMIZE]) ) # Setup distributions. distribution_x = FloatDistribution(low=1.0, high=2.0) distribution_y_1 = CategoricalDistribution(choices=("Shibuya", "Ebisu", "Meguro")) distribution_y_2 = CategoricalDistribution(choices=("Shibuya", "Shinsen")) distribution_z = FloatDistribution(low=1.0, high=100.0, log=True) # Set new params. storage.set_trial_param(trial_id_1, "x", 0.5, distribution_x) storage.set_trial_param(trial_id_1, "y", 2, distribution_y_1) assert storage.get_trial_param(trial_id_1, "x") == 0.5 assert storage.get_trial_param(trial_id_1, "y") == 2 # Check set_param breaks neither get_trial nor get_trial_params. assert storage.get_trial(trial_id_1).params == {"x": 0.5, "y": "Meguro"} assert storage.get_trial_params(trial_id_1) == {"x": 0.5, "y": "Meguro"} # Set params to another trial. storage.set_trial_param(trial_id_2, "x", 0.3, distribution_x) storage.set_trial_param(trial_id_2, "z", 0.1, distribution_z) assert storage.get_trial_param(trial_id_2, "x") == 0.3 assert storage.get_trial_param(trial_id_2, "z") == 0.1 assert storage.get_trial(trial_id_2).params == {"x": 0.3, "z": 0.1} assert storage.get_trial_params(trial_id_2) == {"x": 0.3, "z": 0.1} # Set params with distributions that do not match previous ones. with pytest.raises(ValueError): storage.set_trial_param(trial_id_2, "y", 0.5, distribution_z) # Choices in CategoricalDistribution should match including its order. with pytest.raises(ValueError): storage.set_trial_param( trial_id_2, "y", 2, CategoricalDistribution(choices=("Meguro", "Shibuya", "Ebisu")) ) storage.set_trial_state_values(trial_id_2, state=TrialState.COMPLETE, values=[0.0]) # Cannot assign params to finished trial. with pytest.raises(UpdateFinishedTrialError): storage.set_trial_param(trial_id_2, "y", 2, distribution_y_1) # Check the previous call does not change the params. with pytest.raises(KeyError): storage.get_trial_param(trial_id_2, "y") # State should be checked prior to distribution compatibility. with pytest.raises(UpdateFinishedTrialError): storage.set_trial_param(trial_id_2, "y", 0.4, distribution_z) # Set params of trials in a different study. storage.set_trial_param(trial_id_3, "y", 1, distribution_y_2) assert storage.get_trial_param(trial_id_3, "y") == 1 assert storage.get_trial(trial_id_3).params == {"y": "Shinsen"} assert storage.get_trial_params(trial_id_3) == {"y": "Shinsen"} # Set params of non-existent trial. non_existent_trial_id = max([trial_id_1, trial_id_2, trial_id_3]) + 1 with pytest.raises(KeyError): storage.set_trial_param(non_existent_trial_id, "x", 0.1, distribution_x) def test_set_and_get_trial_param_for_floats(self, storage: BaseStorage) -> None: _test_set_and_get_trial_param_for_floats(storage) def test_set_trial_state_values_for_values(self, storage: BaseStorage) -> None: # Setup test across multiple studies and trials. study_id = storage.create_new_study(directions=[StudyDirection.MINIMIZE]) trial_id_1 = storage.create_new_trial(study_id) trial_id_2 = storage.create_new_trial(study_id) trial_id_3 = storage.create_new_trial( storage.create_new_study(directions=[StudyDirection.MINIMIZE]) ) trial_id_4 = storage.create_new_trial(study_id) trial_id_5 = storage.create_new_trial(study_id) # Test setting new value. storage.set_trial_state_values(trial_id_1, state=TrialState.COMPLETE, values=(0.5,)) storage.set_trial_state_values( trial_id_3, state=TrialState.COMPLETE, values=(float("inf"),) ) storage.set_trial_state_values( trial_id_4, state=TrialState.WAITING, values=(0.1, 0.2, 0.3) ) storage.set_trial_state_values( trial_id_5, state=TrialState.WAITING, values=[0.1, 0.2, 0.3] ) assert storage.get_trial(trial_id_1).value == 0.5 assert storage.get_trial(trial_id_2).value is None assert storage.get_trial(trial_id_3).value == float("inf") assert storage.get_trial(trial_id_4).values == [0.1, 0.2, 0.3] assert storage.get_trial(trial_id_5).values == [0.1, 0.2, 0.3] non_existent_trial_id = max(trial_id_1, trial_id_2, trial_id_3, trial_id_4, trial_id_5) + 1 with pytest.raises(KeyError): storage.set_trial_state_values( non_existent_trial_id, state=TrialState.COMPLETE, values=(1,) ) # Cannot change values of finished trials. with pytest.raises(UpdateFinishedTrialError): storage.set_trial_state_values(trial_id_1, state=TrialState.COMPLETE, values=(1,)) def test_set_trial_intermediate_value(self, storage: BaseStorage) -> None: # Setup test across multiple studies and trials. study_id = storage.create_new_study(directions=[StudyDirection.MINIMIZE]) trial_id_1 = storage.create_new_trial(study_id) trial_id_2 = storage.create_new_trial(study_id) trial_id_3 = storage.create_new_trial( storage.create_new_study(directions=[StudyDirection.MINIMIZE]) ) trial_id_4 = storage.create_new_trial(study_id) # Test setting new values. storage.set_trial_intermediate_value(trial_id_1, 0, 0.3) storage.set_trial_intermediate_value(trial_id_1, 2, 0.4) storage.set_trial_intermediate_value(trial_id_3, 0, 0.1) storage.set_trial_intermediate_value(trial_id_3, 1, 0.4) storage.set_trial_intermediate_value(trial_id_3, 2, 0.5) storage.set_trial_intermediate_value(trial_id_3, 3, float("inf")) storage.set_trial_intermediate_value(trial_id_4, 0, float("nan")) assert storage.get_trial(trial_id_1).intermediate_values == {0: 0.3, 2: 0.4} assert storage.get_trial(trial_id_2).intermediate_values == {} assert storage.get_trial(trial_id_3).intermediate_values == { 0: 0.1, 1: 0.4, 2: 0.5, 3: float("inf"), } assert np.isnan(storage.get_trial(trial_id_4).intermediate_values[0]) # Test setting existing step. storage.set_trial_intermediate_value(trial_id_1, 0, 0.2) assert storage.get_trial(trial_id_1).intermediate_values == {0: 0.2, 2: 0.4} non_existent_trial_id = max(trial_id_1, trial_id_2, trial_id_3, trial_id_4) + 1 with pytest.raises(KeyError): storage.set_trial_intermediate_value(non_existent_trial_id, 0, 0.2) storage.set_trial_state_values(trial_id_1, state=TrialState.COMPLETE, values=[0.0]) # Cannot change values of finished trials. with pytest.raises(UpdateFinishedTrialError): storage.set_trial_intermediate_value(trial_id_1, 0, 0.2) def test_set_trial_intermediate_value_for_floats(self, storage: BaseStorage) -> None: _test_set_trial_intermediate_value_for_floats(storage) def test_get_trial_user_attrs(self, storage: BaseStorage) -> None: _, study_to_trials = _setup_studies(storage, n_study=2, n_trial=5, seed=10) assert all( storage.get_trial_user_attrs(trial_id) == trial.user_attrs for trials in study_to_trials.values() for trial_id, trial in trials.items() ) non_existent_trial = max(tid for ts in study_to_trials.values() for tid in ts.keys()) + 1 with pytest.raises(KeyError): storage.get_trial_user_attrs(non_existent_trial) def test_set_trial_user_attr(self, storage: BaseStorage) -> None: trial_id_1 = storage.create_new_trial( storage.create_new_study(directions=[StudyDirection.MINIMIZE]) ) def check_set_and_get(trial_id: int, key: str, value: Any) -> None: storage.set_trial_user_attr(trial_id, key, value) assert storage.get_trial(trial_id).user_attrs[key] == value # Test setting value. for key, value in EXAMPLE_ATTRS.items(): check_set_and_get(trial_id_1, key, value) assert storage.get_trial(trial_id_1).user_attrs == EXAMPLE_ATTRS # Test overwriting value. check_set_and_get(trial_id_1, "dataset", "ImageNet") # Test another trial. trial_id_2 = storage.create_new_trial( storage.create_new_study(directions=[StudyDirection.MINIMIZE]) ) check_set_and_get(trial_id_2, "baseline_score", 0.001) assert len(storage.get_trial(trial_id_2).user_attrs) == 1 assert storage.get_trial(trial_id_2).user_attrs["baseline_score"] == 0.001 # Cannot set attributes of non-existent trials. non_existent_trial_id = max({trial_id_1, trial_id_2}) + 1 with pytest.raises(KeyError): storage.set_trial_user_attr(non_existent_trial_id, "key", "value") # Cannot set attributes of finished trials. storage.set_trial_state_values(trial_id_1, state=TrialState.COMPLETE, values=[0.0]) with pytest.raises(UpdateFinishedTrialError): storage.set_trial_user_attr(trial_id_1, "key", "value") def test_set_and_get_trial_user_attr_for_floats(self, storage: BaseStorage) -> None: _test_set_and_get_trial_user_attr_for_floats(storage) def test_get_trial_system_attrs(self, storage: BaseStorage) -> None: _, study_to_trials = _setup_studies(storage, n_study=2, n_trial=5, seed=10) assert all( storage.get_trial_system_attrs(trial_id) == trial.system_attrs for trials in study_to_trials.values() for trial_id, trial in trials.items() ) non_existent_trial = max(tid for ts in study_to_trials.values() for tid in ts.keys()) + 1 with pytest.raises(KeyError): storage.get_trial_system_attrs(non_existent_trial) def test_set_trial_system_attr(self, storage: BaseStorage) -> None: study_id = storage.create_new_study(directions=[StudyDirection.MINIMIZE]) trial_id_1 = storage.create_new_trial(study_id) def check_set_and_get(trial_id: int, key: str, value: Any) -> None: storage.set_trial_system_attr(trial_id, key, value) assert storage.get_trial_system_attrs(trial_id)[key] == value # Test setting value. for key, value in EXAMPLE_ATTRS.items(): check_set_and_get(trial_id_1, key, value) system_attrs = storage.get_trial(trial_id_1).system_attrs assert system_attrs == EXAMPLE_ATTRS # Test overwriting value. check_set_and_get(trial_id_1, "dataset", "ImageNet") # Test another trial. trial_id_2 = storage.create_new_trial(study_id) check_set_and_get(trial_id_2, "baseline_score", 0.001) system_attrs = storage.get_trial(trial_id_2).system_attrs assert system_attrs == {"baseline_score": 0.001} # Cannot set attributes of non-existent trials. non_existent_trial_id = max({trial_id_1, trial_id_2}) + 1 with pytest.raises(KeyError): storage.set_trial_system_attr(non_existent_trial_id, "key", "value") # Cannot set attributes of finished trials. storage.set_trial_state_values(trial_id_1, state=TrialState.COMPLETE, values=[0.0]) with pytest.raises(UpdateFinishedTrialError): storage.set_trial_system_attr(trial_id_1, "key", "value") def test_set_and_get_trial_system_attr_for_floats(self, storage: BaseStorage) -> None: _test_set_and_get_trial_system_attr_for_floats(storage) def test_get_all_studies(self, storage: BaseStorage) -> None: expected_frozen_studies, _ = _setup_studies(storage, n_study=10, n_trial=10, seed=46) frozen_studies = storage.get_all_studies() assert len(frozen_studies) == len(expected_frozen_studies) for _, expected_frozen_study in expected_frozen_studies.items(): frozen_study: FrozenStudy | None = None for s in frozen_studies: if s.study_name == expected_frozen_study.study_name: frozen_study = s break assert frozen_study is not None assert frozen_study.direction == expected_frozen_study.direction assert frozen_study.study_name == expected_frozen_study.study_name assert frozen_study.user_attrs == expected_frozen_study.user_attrs assert frozen_study.system_attrs == expected_frozen_study.system_attrs def test_get_trial(self, storage: BaseStorage) -> None: _, study_to_trials = _setup_studies(storage, n_study=2, n_trial=20, seed=47) for _, expected_trials in study_to_trials.items(): for expected_trial in expected_trials.values(): trial = storage.get_trial(expected_trial._trial_id) assert trial == expected_trial non_existent_trial_id = ( max(tid for ts in study_to_trials.values() for tid in ts.keys()) + 1 ) with pytest.raises(KeyError): storage.get_trial(non_existent_trial_id) def test_get_all_trials(self, storage: BaseStorage) -> None: _, study_to_trials = _setup_studies(storage, n_study=2, n_trial=20, seed=48) for study_id, expected_trials in study_to_trials.items(): trials = storage.get_all_trials(study_id) for trial in trials: expected_trial = expected_trials[trial._trial_id] assert trial == expected_trial non_existent_study_id = max(study_to_trials.keys()) + 1 with pytest.raises(KeyError): storage.get_all_trials(non_existent_study_id) @pytest.mark.parametrize("param_names", [["a", "b"], ["b", "a"]]) def test_get_all_trials_params_order( self, storage: BaseStorage, param_names: list[str] ) -> None: study_id = storage.create_new_study(directions=[StudyDirection.MINIMIZE]) trial_id = storage.create_new_trial( study_id, optuna.trial.create_trial(state=TrialState.RUNNING) ) for param_name in param_names: storage.set_trial_param( trial_id, param_name, 1.0, distribution=FloatDistribution(0.0, 2.0) ) trials = storage.get_all_trials(study_id) assert list(trials[0].params.keys()) == param_names def test_get_all_trials_deepcopy_option(self, storage: BaseStorage) -> None: frozen_studies, study_to_trials = _setup_studies(storage, n_study=2, n_trial=5, seed=49) for study_id in frozen_studies: trials0 = storage.get_all_trials(study_id, deepcopy=True) assert len(trials0) == len(study_to_trials[study_id]) # Check modifying output does not break the internal state of the storage. trials0_original = copy.deepcopy(trials0) trials0[0].params["x"] = 0.1 trials1 = storage.get_all_trials(study_id, deepcopy=False) assert trials0_original == trials1 def test_get_all_trials_state_option(self, storage: BaseStorage) -> None: study_id = storage.create_new_study(directions=[StudyDirection.MAXIMIZE]) generator = random.Random(51) states = ( TrialState.COMPLETE, TrialState.COMPLETE, TrialState.PRUNED, ) for state in states: t = _generate_trial(generator) t.state = state t.values = [0.0] if t.state == TrialState.COMPLETE else None storage.create_new_trial(study_id, template_trial=t) trials = storage.get_all_trials(study_id, states=None) assert len(trials) == 3 trials = storage.get_all_trials(study_id, states=(TrialState.COMPLETE,)) assert len(trials) == 2 assert all(t.state == TrialState.COMPLETE for t in trials) trials = storage.get_all_trials(study_id, states=(TrialState.COMPLETE, TrialState.PRUNED)) assert len(trials) == 3 assert all(t.state in (TrialState.COMPLETE, TrialState.PRUNED) for t in trials) trials = storage.get_all_trials(study_id, states=()) assert len(trials) == 0 other_states = [ s for s in ALL_STATES if s != TrialState.COMPLETE and s != TrialState.PRUNED ] for state in other_states: trials = storage.get_all_trials(study_id, states=(state,)) assert len(trials) == 0 def test_get_all_trials_not_modified(self, storage: BaseStorage) -> None: _, study_to_trials = _setup_studies(storage, n_study=2, n_trial=20, seed=48) for study_id in study_to_trials.keys(): trials = storage.get_all_trials(study_id, deepcopy=False) deepcopied_trials = copy.deepcopy(trials) for trial in trials: if not trial.state.is_finished(): storage.set_trial_param(trial._trial_id, "paramX", 0, FloatDistribution(0, 1)) storage.set_trial_user_attr(trial._trial_id, "usr_attrX", 0) storage.set_trial_system_attr(trial._trial_id, "sys_attrX", 0) if trial.state == TrialState.RUNNING: if trial.number % 3 == 0: storage.set_trial_state_values(trial._trial_id, TrialState.COMPLETE, [0]) elif trial.number % 3 == 1: storage.set_trial_intermediate_value(trial._trial_id, 0, 0) storage.set_trial_state_values(trial._trial_id, TrialState.PRUNED, [0]) else: storage.set_trial_state_values(trial._trial_id, TrialState.FAIL) elif trial.state == TrialState.WAITING: storage.set_trial_state_values(trial._trial_id, TrialState.RUNNING) assert trials == deepcopied_trials def test_get_n_trials(self, storage: BaseStorage) -> None: study_id_to_frozen_studies, _ = _setup_studies(storage, n_study=2, n_trial=7, seed=50) for study_id in study_id_to_frozen_studies: assert storage.get_n_trials(study_id) == 7 non_existent_study_id = max(study_id_to_frozen_studies.keys()) + 1 with pytest.raises(KeyError): assert storage.get_n_trials(non_existent_study_id) def test_get_n_trials_state_option(self, storage: BaseStorage) -> None: study_id = storage.create_new_study(directions=(StudyDirection.MAXIMIZE,)) generator = random.Random(51) states = [ TrialState.COMPLETE, TrialState.COMPLETE, TrialState.PRUNED, ] for s in states: t = _generate_trial(generator) t.state = s t.values = [0.0] if t.state == TrialState.COMPLETE else None storage.create_new_trial(study_id, template_trial=t) assert storage.get_n_trials(study_id, TrialState.COMPLETE) == 2 assert storage.get_n_trials(study_id, TrialState.PRUNED) == 1 other_states = [ s for s in ALL_STATES if s != TrialState.COMPLETE and s != TrialState.PRUNED ] for s in other_states: assert storage.get_n_trials(study_id, s) == 0 @pytest.mark.parametrize("direction", [StudyDirection.MAXIMIZE, StudyDirection.MINIMIZE]) @pytest.mark.parametrize( "values", [ [0.0, 1.0, 2.0], [0.0, float("inf"), 1.0], [0.0, float("-inf"), 1.0], [float("inf"), 0.0, 1.0, float("-inf")], [float("inf")], [float("-inf")], ], ) def test_get_best_trial( self, storage: BaseStorage, direction: StudyDirection, values: list[float] ) -> None: study_id = storage.create_new_study(directions=[direction]) with pytest.raises(ValueError): storage.get_best_trial(study_id) with pytest.raises(KeyError): storage.get_best_trial(study_id + 1) generator = random.Random(51) for v in values: template_trial = _generate_trial(generator) template_trial.state = TrialState.COMPLETE template_trial.value = v storage.create_new_trial(study_id, template_trial=template_trial) expected_value = max(values) if direction == StudyDirection.MAXIMIZE else min(values) assert storage.get_best_trial(study_id).value == expected_value def test_get_best_trial_for_multi_objective_optimization(self, storage: BaseStorage) -> None: study_id = storage.create_new_study( directions=(StudyDirection.MAXIMIZE, StudyDirection.MINIMIZE) ) generator = random.Random(51) for i in range(3): template_trial = _generate_trial(generator) template_trial.state = TrialState.COMPLETE template_trial.values = [i, i + 1] storage.create_new_trial(study_id, template_trial=template_trial) with pytest.raises(RuntimeError): storage.get_best_trial(study_id) def test_get_trial_id_from_study_id_trial_number(self, storage: BaseStorage) -> None: with pytest.raises(KeyError): # Matching study does not exist. storage.get_trial_id_from_study_id_trial_number(study_id=0, trial_number=0) study_id = storage.create_new_study(directions=[StudyDirection.MINIMIZE]) with pytest.raises(KeyError): # Matching trial does not exist. storage.get_trial_id_from_study_id_trial_number(study_id, trial_number=0) trial_id = storage.create_new_trial(study_id) assert trial_id == storage.get_trial_id_from_study_id_trial_number( study_id, trial_number=0 ) # Trial IDs are globally unique within a storage but numbers are only unique within a # study. Create a second study within the same storage. study_id = storage.create_new_study(directions=[StudyDirection.MINIMIZE]) trial_id = storage.create_new_trial(study_id) assert trial_id == storage.get_trial_id_from_study_id_trial_number( study_id, trial_number=0 ) def test_pickle_storage(self, storage: BaseStorage) -> None: study_id = storage.create_new_study(directions=[StudyDirection.MINIMIZE]) storage.set_study_system_attr(study_id, "key", "pickle") restored_storage = pickle.loads(pickle.dumps(storage)) storage_system_attrs = storage.get_study_system_attrs(study_id) restored_storage_system_attrs = restored_storage.get_study_system_attrs(study_id) assert storage_system_attrs == restored_storage_system_attrs == {"key": "pickle"} def test_check_trial_is_updatable(self, storage: BaseStorage) -> None: study_id = storage.create_new_study(directions=[StudyDirection.MINIMIZE]) trial_id = storage.create_new_trial(study_id) storage.check_trial_is_updatable(trial_id, TrialState.RUNNING) storage.check_trial_is_updatable(trial_id, TrialState.WAITING) with pytest.raises(UpdateFinishedTrialError): storage.check_trial_is_updatable(trial_id, TrialState.FAIL) with pytest.raises(UpdateFinishedTrialError): storage.check_trial_is_updatable(trial_id, TrialState.PRUNED) with pytest.raises(UpdateFinishedTrialError): storage.check_trial_is_updatable(trial_id, TrialState.COMPLETE) ALL_STATES = list(TrialState) EXAMPLE_ATTRS: dict[str, JSONSerializable] = { "dataset": "MNIST", "none": None, "json_serializable": {"baseline_score": 0.001, "tags": ["image", "classification"]}, } FLOAT_ATTRS = { "zero": 0, "pi": math.pi, "max": sys.float_info.max, "negative max": -sys.float_info.max, "min": sys.float_info.min, "negative min": -sys.float_info.min, "inf": float("inf"), "negative inf": -float("inf"), "nan": float("nan"), } def is_equal_floats(a: float, b: float) -> bool: if math.isnan(a): return math.isnan(b) if math.isnan(b): return False return a == b def _test_set_and_get_study_user_attrs_for_floats(storage: BaseStorage) -> None: study_id = storage.create_new_study(directions=[StudyDirection.MINIMIZE]) # Test setting value. for key, value in FLOAT_ATTRS.items(): storage.set_study_user_attr(study_id, key, value) assert is_equal_floats(storage.get_study_user_attrs(study_id)[key], value) def _test_set_and_get_study_system_attrs_for_floats(storage: BaseStorage) -> None: study_id = storage.create_new_study(directions=[StudyDirection.MINIMIZE]) # Test setting value. for key, value in FLOAT_ATTRS.items(): storage.set_study_system_attr(study_id, key, value) assert is_equal_floats(storage.get_study_system_attrs(study_id)[key], value) def _test_set_trial_state_values_for_floats(storage: BaseStorage) -> None: study_id = storage.create_new_study(directions=[StudyDirection.MINIMIZE]) for value in FLOAT_ATTRS.values(): if math.isnan(value): # NOTE: Optuna does not accept `nan` as `value`. continue trial_id = storage.create_new_trial(study_id) storage.set_trial_state_values(trial_id, state=TrialState.COMPLETE, values=(value,)) set_value = storage.get_trial(trial_id).value assert set_value is not None assert is_equal_floats(set_value, value) def _test_set_and_get_trial_param_for_floats(storage: BaseStorage) -> None: study_id = storage.create_new_study(directions=[StudyDirection.MINIMIZE]) trial_id = storage.create_new_trial(study_id) for key, value in FLOAT_ATTRS.items(): float_distribution = FloatDistribution(low=value, high=value) categorical_distribution = CategoricalDistribution(choices=(value,)) for distribution in (float_distribution, categorical_distribution): # NOTE: suggest_float does not generate infinite values and NaN. if isinstance(distribution, FloatDistribution) and not math.isfinite(value): continue param_name = distribution.__class__.__name__ + key internal_repr = distribution.to_internal_repr(value) storage.set_trial_param(trial_id, param_name, internal_repr, distribution) assert is_equal_floats(storage.get_trial_param(trial_id, param_name), internal_repr) assert storage.get_trial(trial_id).distributions[param_name] == distribution def _test_set_trial_intermediate_value_for_floats(storage: BaseStorage) -> None: study_id = storage.create_new_study(directions=[StudyDirection.MINIMIZE]) trial_id = storage.create_new_trial(study_id) for i, value in enumerate(FLOAT_ATTRS.values()): storage.set_trial_intermediate_value(trial_id, i, value) assert is_equal_floats(storage.get_trial(trial_id).intermediate_values[i], value) def _test_set_and_get_trial_user_attr_for_floats(storage: BaseStorage) -> None: trial_id = storage.create_new_trial( storage.create_new_study(directions=[StudyDirection.MINIMIZE]) ) # Test setting value. for key, value in FLOAT_ATTRS.items(): storage.set_trial_user_attr(trial_id, key, value) assert is_equal_floats(storage.get_trial_user_attrs(trial_id)[key], value) def _test_set_and_get_trial_system_attr_for_floats(storage: BaseStorage) -> None: trial_id = storage.create_new_trial( storage.create_new_study(directions=[StudyDirection.MINIMIZE]) ) # Test setting value. for key, value in FLOAT_ATTRS.items(): storage.set_trial_system_attr(trial_id, key, value) assert is_equal_floats(storage.get_trial_system_attrs(trial_id)[key], value) def _setup_studies( storage: BaseStorage, n_study: int, n_trial: int, seed: int, direction: StudyDirection | None = None, ) -> tuple[dict[int, FrozenStudy], dict[int, dict[int, FrozenTrial]]]: generator = random.Random(seed) study_id_to_frozen_study: dict[int, FrozenStudy] = {} study_id_to_trials: dict[int, dict[int, FrozenTrial]] = {} for i in range(n_study): study_name = f"test-study-name-{i}" if direction is None: direction = generator.choice([StudyDirection.MINIMIZE, StudyDirection.MAXIMIZE]) study_id = storage.create_new_study(directions=(direction,), study_name=study_name) storage.set_study_user_attr(study_id, "u", i) storage.set_study_system_attr(study_id, "s", i) trials = {} for j in range(n_trial): trial = _generate_trial(generator) trial.number = j trial._trial_id = storage.create_new_trial(study_id, trial) trials[trial._trial_id] = trial study_id_to_trials[study_id] = trials study_id_to_frozen_study[study_id] = FrozenStudy( study_name=study_name, direction=direction, user_attrs={"u": i}, system_attrs={"s": i}, study_id=study_id, ) return study_id_to_frozen_study, study_id_to_trials def _generate_trial(generator: random.Random) -> FrozenTrial: example_params = { "paramA": (generator.uniform(0, 1), FloatDistribution(0, 1)), "paramB": (generator.uniform(1, 2), FloatDistribution(1, 2, log=True)), "paramC": ( generator.choice(["CatA", "CatB", "CatC"]), CategoricalDistribution(("CatA", "CatB", "CatC")), ), "paramD": (generator.uniform(-3, 0), FloatDistribution(-3, 0)), "paramE": (generator.choice([0.1, 0.2]), CategoricalDistribution((0.1, 0.2))), } example_attrs = { "attrA": "valueA", "attrB": 1, "attrC": None, "attrD": {"baseline_score": 0.001, "tags": ["image", "classification"]}, } state = generator.choice(ALL_STATES) params = {} distributions = {} user_attrs = {} system_attrs: dict[str, Any] = {} intermediate_values = {} for key, (value, dist) in example_params.items(): if generator.choice([True, False]): params[key] = value distributions[key] = dist for key, value in example_attrs.items(): if generator.choice([True, False]): user_attrs["usr_" + key] = value if generator.choice([True, False]): system_attrs["sys_" + key] = value for i in range(generator.randint(4, 10)): if generator.choice([True, False]): intermediate_values[i] = generator.uniform(-10, 10) return FrozenTrial( number=0, # dummy state=state, value=generator.uniform(-10, 10) if state == TrialState.COMPLETE else None, datetime_start=datetime.now(), datetime_complete=datetime.now() if state.is_finished() else None, params=params, distributions=distributions, user_attrs=user_attrs, system_attrs=system_attrs, intermediate_values=intermediate_values, trial_id=0, # dummy )