Add optuna_dashboard.preferential module

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",
+]

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)

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]

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

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

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()