Merge pull request #546 from c-bata/preferential-optimization

Add `optuna_dashboard.preferential` module

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