diff --git a/optuna_dashboard/preferential/__init__.py b/optuna_dashboard/preferential/__init__.py new file mode 100644 index 00000000..7760d32f --- /dev/null +++ b/optuna_dashboard/preferential/__init__.py @@ -0,0 +1,12 @@ +from __future__ import annotations + +from optuna_dashboard.preferential._study import create_study +from optuna_dashboard.preferential._study import load_study +from optuna_dashboard.preferential._study import PreferentialStudy + + +__all__ = [ + "PreferentialStudy", + "create_study", + "load_study", +] diff --git a/optuna_dashboard/preferential/_study.py b/optuna_dashboard/preferential/_study.py new file mode 100644 index 00000000..9bbd6d2f --- /dev/null +++ b/optuna_dashboard/preferential/_study.py @@ -0,0 +1,152 @@ +from __future__ import annotations + +import copy +from typing import Any +from typing import Container +from typing import Iterable + +import optuna +from optuna import logging +from optuna.distributions import BaseDistribution +from optuna.samplers import BaseSampler +from optuna.samplers import RandomSampler +from optuna.trial import FrozenTrial +from optuna.trial import TrialState +from optuna_dashboard.preferential._system_attrs import get_preferences +from optuna_dashboard.preferential._system_attrs import report_preferences + + +_logger = logging.get_logger(__name__) +_SYSTEM_ATTR_PREFERENTIAL_STUDY = "preference:is_preferential" +_SYSTEM_ATTR_COMPARISON_READY = "preference:comparison_ready" + + +class PreferentialStudy: + def __init__(self, study: optuna.Study) -> None: + self._study = study + + @property + def trials(self) -> list[FrozenTrial]: + return self._study.trials + + @property + def best_trials(self) -> list[FrozenTrial]: + ready_trials = [ + t + for t in self._study.get_trials( + deepcopy=False, states=(TrialState.COMPLETE, TrialState.RUNNING) + ) + if t.system_attrs.get(_SYSTEM_ATTR_COMPARISON_READY) is True + ] + preferences = get_preferences(self._study, deepcopy=False) + worse_numbers = {worse.number for _, worse in preferences} + return [copy.deepcopy(t) for t in ready_trials if t.number not in worse_numbers] + + @property + def study_name(self) -> str: + return self._study.study_name + + @property + def user_attrs(self) -> dict[str, Any]: + return self._study.user_attrs + + @property + def preferences(self) -> list[tuple[FrozenTrial, FrozenTrial]]: + return self.get_preferences(deepcopy=True) + + def get_trials( + self, + deepcopy: bool = True, + states: Container[optuna.trial.TrialState] | None = None, + ) -> list[FrozenTrial]: + return self._study.get_trials(deepcopy, states) + + def ask(self, fixed_distributions: dict[str, BaseDistribution] | None = None) -> optuna.Trial: + return self._study.ask(fixed_distributions) + + def add_trial(self, trial: FrozenTrial) -> None: + self._study.add_trial(trial) + + def add_trials(self, trials: Iterable[FrozenTrial]) -> None: + self._study.add_trials(trials) + + def report_preference( + self, + better_trials: FrozenTrial | list[FrozenTrial], + worse_trials: FrozenTrial | list[FrozenTrial], + ) -> None: + if not isinstance(better_trials, list): + better_trials = [better_trials] + if not isinstance(worse_trials, list): + worse_trials = [worse_trials] + + report_preferences(self._study, [(b, w) for b in better_trials for w in worse_trials]) + + def get_preferences(self, *, deepcopy: bool = True) -> list[tuple[FrozenTrial, FrozenTrial]]: + return get_preferences(self._study, deepcopy=deepcopy) + + def set_user_attr(self, key: str, value: Any) -> None: + self._study.set_user_attr(key, value) + + def mark_comparison_ready(self, trial_or_number: optuna.Trial | int) -> None: + storage = self._study._storage + if isinstance(trial_or_number, optuna.Trial): + trial_id = trial_or_number._trial_id + elif isinstance(trial_or_number, int): + trial_id = storage.get_trial_id_from_study_id_trial_number( + self._study._study_id, trial_or_number + ) + else: + raise RuntimeError("Unexpected trial type") + storage.set_trial_system_attr(trial_id, _SYSTEM_ATTR_COMPARISON_READY, True) + + +def create_study( + *, + storage: str | optuna.storages.BaseStorage | None = None, + sampler: BaseSampler | None = None, + study_name: str | None = None, + load_if_exists: bool = False, +) -> PreferentialStudy: + try: + study = optuna.create_study( + storage=storage, + sampler=sampler or RandomSampler(), + study_name=study_name, + ) + study._storage.set_study_system_attr( + study._study_id, _SYSTEM_ATTR_PREFERENTIAL_STUDY, True + ) + return PreferentialStudy(study) + + except optuna.exceptions.DuplicatedStudyError: + if load_if_exists: + assert study_name is not None + assert storage is not None + + _logger.info( + "Using an existing study with name '{}' instead of " + "creating a new one.".format(study_name) + ) + return load_study( + study_name=study_name, + storage=storage, + sampler=sampler, + ) + else: + raise + + +def load_study( + *, + study_name: str | None, + storage: str | optuna.storages.BaseStorage, + sampler: BaseSampler | None = None, +) -> PreferentialStudy: + study = optuna.load_study( + study_name=study_name, storage=storage, sampler=sampler or RandomSampler() + ) + system_attrs = study._storage.get_study_system_attrs(study._study_id) + if not system_attrs.get(_SYSTEM_ATTR_PREFERENTIAL_STUDY): + raise ValueError("The study is not a PreferentialStudy.") + return PreferentialStudy(study) diff --git a/optuna_dashboard/preferential/_system_attrs.py b/optuna_dashboard/preferential/_system_attrs.py new file mode 100644 index 00000000..70567655 --- /dev/null +++ b/optuna_dashboard/preferential/_system_attrs.py @@ -0,0 +1,46 @@ +from __future__ import annotations + +import uuid + +import optuna +from optuna.trial import FrozenTrial +from optuna.trial import TrialState + + +_SYSTEM_ATTR_PREFIX_PREFERENCE = "preference:values" + + +def report_preferences( + study: optuna.Study, + preferences: list[tuple[FrozenTrial, FrozenTrial]], +) -> None: + key = _SYSTEM_ATTR_PREFIX_PREFERENCE + str(uuid.uuid4()) + study._storage.set_study_system_attr( + study_id=study._study_id, + key=key, + value=[(better.number, worse.number) for better, worse in preferences], + ) + + values = [0 for _ in study.directions] + for better, worse in preferences: + for t in (better, worse): + study.tell( + t.number, + values=values, + state=TrialState.COMPLETE, + skip_if_finished=True, + ) + + +def get_preferences( + study: optuna.Study, + *, + deepcopy: bool = True, +) -> list[tuple[FrozenTrial, FrozenTrial]]: + preferences: list[tuple[int, int]] = [] + for k, v in study.system_attrs.items(): + if not k.startswith(_SYSTEM_ATTR_PREFIX_PREFERENCE): + continue + preferences.extend(v) # type: ignore + trials = study.get_trials(deepcopy=deepcopy) + return [(trials[better], trials[worse]) for (better, worse) in preferences] diff --git a/python_tests/preferential/__init__.py b/python_tests/preferential/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/python_tests/preferential/test_study.py b/python_tests/preferential/test_study.py new file mode 100644 index 00000000..3e0b011c --- /dev/null +++ b/python_tests/preferential/test_study.py @@ -0,0 +1,326 @@ +from __future__ import annotations + +import copy +import multiprocessing +import pickle +from typing import Callable +from unittest.mock import patch +import uuid + +from optuna import copy_study +from optuna import create_trial +from optuna import delete_study +from optuna import distributions +from optuna import Trial +from optuna.exceptions import DuplicatedStudyError +from optuna.trial import TrialState +from optuna_dashboard.preferential import create_study +from optuna_dashboard.preferential import load_study +import pytest + +from ..storage_supplier import parametrize_storages +from ..storage_supplier import StorageSupplier + + +@parametrize_storages +def test_study_set_and_get_user_attrs(storage_supplier: Callable[[], StorageSupplier]) -> None: + with storage_supplier() as storage: + study = create_study(storage=storage) + + study.set_user_attr("dataset", "MNIST") + assert study.user_attrs["dataset"] == "MNIST" + + +@parametrize_storages +def test_report_and_get_preferences(storage_supplier: Callable[[], StorageSupplier]) -> None: + with storage_supplier() as storage: + study = create_study(storage=storage) + assert len(study.preferences) == 0 + + for _ in range(2): + trial = study.ask() + trial.suggest_float("x", 0, 1) + study.mark_comparison_ready(trial) + better, worse = study.trials + study.report_preference(better, worse) + assert len(study.preferences) == 1 + + actual_better, actual_worse = study.preferences[0] + assert actual_better.number == better.number + assert actual_worse.number == worse.number + + +def test_study_pickle() -> None: + study_1 = create_study() + for _ in range(10): + study_1.ask() + assert len(study_1.trials) == 10 + dumped_bytes = pickle.dumps(study_1) + + study_2 = pickle.loads(dumped_bytes) + assert len(study_2.trials) == 10 + + for _ in range(10): + study_2.ask() + assert len(study_2.trials) == 20 + + +@parametrize_storages +def test_create_study(storage_supplier: Callable[[], StorageSupplier]) -> None: + with storage_supplier() as storage: + # Test creating a new study. + study = create_study(storage=storage, load_if_exists=False) + + # Test `load_if_exists=True` with existing study. + create_study(study_name=study.study_name, storage=storage, load_if_exists=True) + + with pytest.raises(DuplicatedStudyError): + create_study(study_name=study.study_name, storage=storage, load_if_exists=False) + + +@parametrize_storages +def test_load_study(storage_supplier: Callable[[], StorageSupplier]) -> None: + with storage_supplier() as storage: + if storage is None: + # `InMemoryStorage` can not be used with `load_study` function. + return + + study_name = str(uuid.uuid4()) + + with pytest.raises(KeyError): + # Test loading an unexisting study. + load_study(study_name=study_name, storage=storage) + + # Create a new study. + created_study = create_study(study_name=study_name, storage=storage) + + # Test loading an existing study. + loaded_study = load_study(study_name=study_name, storage=storage) + assert created_study.study_name == loaded_study.study_name + + +@parametrize_storages +def test_load_study_study_name_none(storage_supplier: Callable[[], StorageSupplier]) -> None: + with storage_supplier() as storage: + if storage is None: + # `InMemoryStorage` can not be used with `load_study` function. + return + + study_name = str(uuid.uuid4()) + + _ = create_study(study_name=study_name, storage=storage) + + loaded_study = load_study(study_name=None, storage=storage) + + assert loaded_study.study_name == study_name + + study_name = str(uuid.uuid4()) + + _ = create_study(study_name=study_name, storage=storage) + + # Ambiguous study. + with pytest.raises(ValueError): + load_study(study_name=None, storage=storage) + + +@parametrize_storages +def test_delete_study(storage_supplier: Callable[[], StorageSupplier]) -> None: + with storage_supplier() as storage: + # Test deleting a non-existing study. + with pytest.raises(KeyError): + delete_study(study_name="invalid-study-name", storage=storage) + + # Test deleting an existing study. + study = create_study(storage=storage, load_if_exists=False) + delete_study(study_name=study.study_name, storage=storage) + + # Test failed to delete the study which is already deleted. + with pytest.raises(KeyError): + delete_study(study_name=study.study_name, storage=storage) + + +def test_copy_study() -> None: + with StorageSupplier("sqlite") as from_storage, StorageSupplier("sqlite") as to_storage: + from_study = create_study(storage=from_storage) + from_study.set_user_attr("baz", "qux") + for _ in range(3): + trial = from_study.ask() + trial.suggest_float("x", 0, 1) + from_study.mark_comparison_ready(trial) + from_study.report_preference(from_study.trials[0], from_study.trials[1]) + from_study.report_preference(from_study.trials[1], from_study.trials[2]) + + copy_study( + from_study_name=from_study.study_name, + from_storage=from_storage, + to_storage=to_storage, + ) + + to_study = load_study(study_name=from_study.study_name, storage=to_storage) + assert to_study.study_name == from_study.study_name + assert to_study.user_attrs == from_study.user_attrs + assert len(to_study.trials) == len(from_study.trials) + assert len(from_study.preferences) == len(to_study.preferences) + + +def test_copy_study_to_study_name() -> None: + with StorageSupplier("sqlite") as from_storage, StorageSupplier("sqlite") as to_storage: + from_study = create_study(study_name="foo", storage=from_storage) + _ = create_study(study_name="foo", storage=to_storage) + + with pytest.raises(DuplicatedStudyError): + copy_study( + from_study_name=from_study.study_name, + from_storage=from_storage, + to_storage=to_storage, + ) + + copy_study( + from_study_name=from_study.study_name, + from_storage=from_storage, + to_storage=to_storage, + to_study_name="bar", + ) + + _ = load_study(study_name="bar", storage=to_storage) + + +@parametrize_storages +def test_add_trial(storage_supplier: Callable[[], StorageSupplier]) -> None: + with storage_supplier() as storage: + study = create_study(storage=storage) + assert len(study.trials) == 0 + + trial = create_trial(value=0) + study.add_trial(trial) + assert len(study.trials) == 1 + assert study.trials[0].number == 0 + + +def test_add_trial_invalid_values_length() -> None: + study = create_study() + trial = create_trial(values=[0, 0]) + with pytest.raises(ValueError): + study.add_trial(trial) + + +@parametrize_storages +def test_add_trials(storage_supplier: Callable[[], StorageSupplier]) -> None: + with storage_supplier() as storage: + study = create_study(storage=storage) + assert len(study.trials) == 0 + + study.add_trials([]) + assert len(study.trials) == 0 + + trials = [create_trial(value=i) for i in range(3)] + study.add_trials(trials) + assert len(study.trials) == 3 + for i, trial in enumerate(study.trials): + assert trial.number == i + assert trial.value == i + + other_study = create_study(storage=storage) + other_study.add_trials(study.trials) + assert len(other_study.trials) == 3 + for i, trial in enumerate(other_study.trials): + assert trial.number == i + assert trial.value == i + + +@parametrize_storages +def test_get_trials(storage_supplier: Callable[[], StorageSupplier]) -> None: + with storage_supplier() as storage: + study = create_study(storage=storage) + for _ in range(5): + trial = study.ask() + trial.suggest_int("x", 1, 5) + study.mark_comparison_ready(trial) + + with patch("copy.deepcopy", wraps=copy.deepcopy) as mock_object: + trials0 = study.get_trials(deepcopy=False) + assert mock_object.call_count == 0 + assert len(trials0) == 5 + + trials1 = study.get_trials(deepcopy=True) + assert mock_object.call_count > 0 + assert trials0 == trials1 + + # `study.trials` is equivalent to `study.get_trials(deepcopy=True)`. + old_count = mock_object.call_count + trials2 = study.trials + assert mock_object.call_count > old_count + assert trials0 == trials2 + + +@parametrize_storages +def test_get_trials_state_option(storage_supplier: Callable[[], StorageSupplier]) -> None: + with storage_supplier() as storage: + study = create_study(storage=storage) + for _ in range(3): + trial = study.ask() + study.mark_comparison_ready(trial) + better, worse = study.trials[:2] + study.report_preference(better, worse) + + trials = study.get_trials(states=None) + assert len(trials) == 3 + + trials = study.get_trials(states=(TrialState.RUNNING,)) + assert len(trials) == 1 + assert all(t.state == TrialState.RUNNING for t in trials) + + trials = study.get_trials(states=(TrialState.COMPLETE,)) + assert len(trials) == 2 + assert all(t.state == TrialState.COMPLETE for t in trials) + + trials = study.get_trials(states=()) + assert len(trials) == 0 + + other_states = [ + s for s in list(TrialState) if s != TrialState.COMPLETE and s != TrialState.RUNNING + ] + for s in other_states: + trials = study.get_trials(states=(s,)) + assert len(trials) == 0 + + +def test_ask() -> None: + study = create_study() + + trial = study.ask() + assert isinstance(trial, Trial) + + +def test_ask_fixed_search_space() -> None: + fixed_distributions = { + "x": distributions.FloatDistribution(0, 1), + "y": distributions.CategoricalDistribution(["bacon", "spam"]), + } + + study = create_study() + trial = study.ask(fixed_distributions=fixed_distributions) + + params = trial.params + assert len(trial.params) == 2 + assert 0 <= params["x"] < 1 + assert params["y"] in ["bacon", "spam"] + + +def test_report_preferences_from_another_process() -> None: + pool = multiprocessing.Pool() + + with StorageSupplier("sqlite") as storage: + # Create a study and ask for a new trial. + study = create_study(storage=storage) + study.ask() + study.ask() + + # Test normal behaviour. + better, worse = study.trials + pool.starmap(study.report_preference, [(better, worse)]) + + assert len(study.trials) == 2 + assert study.trials[0].state == TrialState.COMPLETE + assert study.trials[1].state == TrialState.COMPLETE + assert len(study.preferences) == 1 diff --git a/python_tests/preferential/test_system_attrs.py b/python_tests/preferential/test_system_attrs.py new file mode 100644 index 00000000..b9e22e19 --- /dev/null +++ b/python_tests/preferential/test_system_attrs.py @@ -0,0 +1,28 @@ +from __future__ import annotations + +from typing import Callable + +import optuna +from optuna_dashboard.preferential._system_attrs import get_preferences +from optuna_dashboard.preferential._system_attrs import report_preferences + +from ..storage_supplier import parametrize_storages +from ..storage_supplier import StorageSupplier + + +@parametrize_storages +def test_report_and_get_preferences(storage_supplier: Callable[[], StorageSupplier]) -> None: + with storage_supplier() as storage: + study = optuna.create_study(storage=storage) + study.ask() + study.ask() + + assert len(get_preferences(study)) == 0 + + better, worse = study.trials[0], study.trials[1] + report_preferences(study, [(better, worse)]) + assert len(get_preferences(study)) == 1 + + actual_better, actual_worse = get_preferences(study)[0] + assert actual_better.number == better.number + assert actual_worse.number == worse.number diff --git a/python_tests/storage_supplier.py b/python_tests/storage_supplier.py new file mode 100644 index 00000000..8fa53ea8 --- /dev/null +++ b/python_tests/storage_supplier.py @@ -0,0 +1,75 @@ +from __future__ import annotations + +import tempfile +from types import TracebackType +from typing import Any +from typing import IO + +import optuna +from optuna.version import __version__ as optuna_ver +from packaging import version +import pytest + + +parametrize_storages = pytest.mark.parametrize( + "storage_supplier", + [ + lambda: StorageSupplier("inmemory"), + lambda: StorageSupplier("sqlite"), + lambda: StorageSupplier("cached_sqlite"), + # TODO(c-bata): Support "JournalRedisStorage" + pytest.param( + lambda: StorageSupplier("journal"), + marks=pytest.mark.skipif( + version.parse(optuna_ver) < version.Version("3.1.0"), + reason="Artifact is not implemented yet in Optuna", + ), + ), + ], +) +SQLITE3_TIMEOUT = 300 + + +class StorageSupplier: + def __init__(self, storage_specifier: str, **kwargs: Any) -> None: + self.storage_specifier = storage_specifier + self.extra_args = kwargs + self.tempfile: IO[Any] | None = None + + def __enter__( + self, + ) -> ( + optuna.storages.InMemoryStorage + | optuna.storages._CachedStorage + | optuna.storages.RDBStorage + | optuna.storages.JournalStorage + ): + if self.storage_specifier == "inmemory": + if len(self.extra_args) > 0: + raise ValueError("InMemoryStorage does not accept any arguments!") + return optuna.storages.InMemoryStorage() + elif "sqlite" in self.storage_specifier: + self.tempfile = tempfile.NamedTemporaryFile(**self.extra_args) + url = "sqlite:///{}".format(self.tempfile.name) + rdb_storage = optuna.storages.RDBStorage( + url, + engine_kwargs={"connect_args": {"timeout": SQLITE3_TIMEOUT}}, + **self.extra_args, + ) + return ( + optuna.storages._CachedStorage(rdb_storage) + if "cached" in self.storage_specifier + else rdb_storage + ) + elif "journal" in self.storage_specifier: + self.tempfile = tempfile.NamedTemporaryFile(**self.extra_args) + file_storage = optuna.storages.JournalFileStorage(self.tempfile.name) + return optuna.storages.JournalStorage(file_storage) + else: + assert False, "Must not reach here" + + def __exit__( + self, exc_type: type[BaseException], exc_val: BaseException, exc_tb: TracebackType + ) -> None: + if self.tempfile: + self.tempfile.close()