diff --git a/optuna_dashboard/_app.py b/optuna_dashboard/_app.py index ab26a1a6..7aae8864 100644 --- a/optuna_dashboard/_app.py +++ b/optuna_dashboard/_app.py @@ -38,6 +38,9 @@ from ._storage_url import get_storage from .artifact._backend import delete_all_artifacts from .artifact._backend import register_artifact_route from .artifact._backend_to_store import to_artifact_store +from .preferential._study import _SYSTEM_ATTR_PREFERENTIAL_STUDY +from .preferential._study import get_best_trials as get_best_preferential_trials +from .preferential._system_attrs import report_preferences if typing.TYPE_CHECKING: @@ -187,8 +190,12 @@ def create_app( return {"reason": f"study_id={study_id} is not found"} trials = get_trials(storage, study_id) + system_attrs = getattr(summary, "system_attrs", {}) + is_preferential = system_attrs.get(_SYSTEM_ATTR_PREFERENTIAL_STUDY, False) # TODO(c-bata): Cache best_trials - if len(summary.directions) == 1: + if is_preferential: + best_trials = get_best_preferential_trials(study_id, storage) + elif len(summary.directions) == 1: if len([t for t in trials if t.state == TrialState.COMPLETE]) == 0: best_trials = [] else: @@ -255,6 +262,25 @@ def create_app( response.status = 204 # No content return {} + @app.post("/api/studies//preference") + @json_api_view + def post_preference(study_id: int) -> dict[str, Any]: + try: + best_trials = [int(d) for d in request.json.get("best_trials", [])] + worst_trials = [int(d) for d in request.json.get("worst_trials", [])] + except ValueError: + response.status = 400 + return {"reason": "best_trials and worst_trials must be an array of integers."} + if len(best_trials) == 0 or len(worst_trials) == 0: + response.status = 400 # Bad request + return {"reason": "You need to set best_trials and worst_trials"} + + preferences = [(best, worst) for best in best_trials for worst in worst_trials] + report_preferences(study_id, storage, preferences) + + response.status = 204 + return {} + @app.post("/api/trials//tell") @json_api_view def tell_trial(trial_id: int) -> dict[str, Any]: diff --git a/optuna_dashboard/preferential/_study.py b/optuna_dashboard/preferential/_study.py index 9bbd6d2f..71c00882 100644 --- a/optuna_dashboard/preferential/_study.py +++ b/optuna_dashboard/preferential/_study.py @@ -31,16 +31,7 @@ class PreferentialStudy: @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] + return get_best_trials(self._study._study_id, self._study._storage) @property def study_name(self) -> str: @@ -80,10 +71,16 @@ class PreferentialStudy: 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]) + report_preferences( + self._study._study_id, + self._study._storage, + [(b.number, w.number) 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) + trials = self._study.get_trials(deepcopy=deepcopy) + preferences = get_preferences(self._study._study_id, self._study._storage) + return [(trials[better], trials[worse]) for (better, worse) in preferences] def set_user_attr(self, key: str, value: Any) -> None: self._study.set_user_attr(key, value) @@ -101,6 +98,21 @@ class PreferentialStudy: storage.set_trial_system_attr(trial_id, _SYSTEM_ATTR_COMPARISON_READY, True) +def get_best_trials(study_id: int, storage: optuna.storages.BaseStorage) -> list[FrozenTrial]: + ready_trials = [ + t + for t in storage.get_all_trials( + study_id, + deepcopy=False, + states=(TrialState.COMPLETE, TrialState.RUNNING), + ) + if t.system_attrs.get(_SYSTEM_ATTR_COMPARISON_READY) is True + ] + preferences = get_preferences(study_id, storage) + worse_numbers = {worse for _, worse in preferences} + return [copy.deepcopy(t) for t in ready_trials if t.number not in worse_numbers] + + def create_study( *, storage: str | optuna.storages.BaseStorage | None = None, diff --git a/optuna_dashboard/preferential/_system_attrs.py b/optuna_dashboard/preferential/_system_attrs.py index 70567655..fdd9db35 100644 --- a/optuna_dashboard/preferential/_system_attrs.py +++ b/optuna_dashboard/preferential/_system_attrs.py @@ -2,45 +2,45 @@ from __future__ import annotations import uuid -import optuna -from optuna.trial import FrozenTrial +from optuna.storages import BaseStorage from optuna.trial import TrialState +from .._storage import get_study_summary + _SYSTEM_ATTR_PREFIX_PREFERENCE = "preference:values" def report_preferences( - study: optuna.Study, - preferences: list[tuple[FrozenTrial, FrozenTrial]], + study_id: int, + storage: BaseStorage, + preferences: list[tuple[int, int]], ) -> None: key = _SYSTEM_ATTR_PREFIX_PREFERENCE + str(uuid.uuid4()) - study._storage.set_study_system_attr( - study_id=study._study_id, + storage.set_study_system_attr( + study_id=study_id, key=key, - value=[(better.number, worse.number) for better, worse in preferences], + value=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, - ) + trials = storage.get_all_trials(study_id, deepcopy=False) + directions = storage.get_study_directions(study_id) + values = [0 for _ in directions] + updated_trials = {num for tpl in preferences for num in tpl} + for number in updated_trials: + trial_id = trials[number]._trial_id + if trials[number].state != TrialState.COMPLETE: + storage.set_trial_state_values(trial_id, TrialState.COMPLETE, values) def get_preferences( - study: optuna.Study, - *, - deepcopy: bool = True, -) -> list[tuple[FrozenTrial, FrozenTrial]]: + study_id: int, + storage: BaseStorage, +) -> list[tuple[int, int]]: preferences: list[tuple[int, int]] = [] - for k, v in study.system_attrs.items(): + summary = get_study_summary(storage, study_id) + system_attrs = getattr(summary, "system_attrs", {}) + for k, v in 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] + return preferences diff --git a/python_tests/preferential/test_system_attrs.py b/python_tests/preferential/test_system_attrs.py index b9e22e19..10448d48 100644 --- a/python_tests/preferential/test_system_attrs.py +++ b/python_tests/preferential/test_system_attrs.py @@ -17,12 +17,13 @@ def test_report_and_get_preferences(storage_supplier: Callable[[], StorageSuppli study.ask() study.ask() - assert len(get_preferences(study)) == 0 + study_id = study._study_id + assert len(get_preferences(study_id, storage)) == 0 better, worse = study.trials[0], study.trials[1] - report_preferences(study, [(better, worse)]) - assert len(get_preferences(study)) == 1 + report_preferences(study_id, storage, [(better.number, worse.number)]) + assert len(get_preferences(study_id, storage)) == 1 - actual_better, actual_worse = get_preferences(study)[0] - assert actual_better.number == better.number - assert actual_worse.number == worse.number + actual_better, actual_worse = get_preferences(study_id, storage)[0] + assert actual_better == better.number + assert actual_worse == worse.number diff --git a/python_tests/test_api.py b/python_tests/test_api.py index 4dd2b3d2..c995c0e2 100644 --- a/python_tests/test_api.py +++ b/python_tests/test_api.py @@ -8,6 +8,7 @@ from optuna import get_all_study_summaries from optuna.study import StudyDirection from optuna_dashboard._app import create_app from optuna_dashboard._app import create_new_study +from optuna_dashboard.preferential import create_study from .wsgi_client import send_request @@ -99,6 +100,57 @@ class APITestCase(TestCase): ) self.assertEqual(status, 400) + def test_get_best_trials_of_preferential_study(self) -> None: + storage = optuna.storages.InMemoryStorage() + study = create_study(storage=storage) + for _ in range(3): + trial = study.ask() + study.mark_comparison_ready(trial) + study.report_preference(study.trials[0], study.trials[1]) + + app = create_app(storage) + study_id = study._study._study_id + status, _, body = send_request( + app, + f"/api/studies/{study_id}", + "GET", + content_type="application/json", + ) + self.assertEqual(status, 200) + + best_trials = json.loads(body)["best_trials"] + assert len(best_trials) == 2 + assert best_trials[0]["number"] == 0 + assert best_trials[1]["number"] == 2 + + def test_report_preference(self) -> None: + storage = optuna.storages.InMemoryStorage() + study = create_study(storage=storage) + for _ in range(3): + trial = study.ask() + study.mark_comparison_ready(trial) + + app = create_app(storage) + study_id = study._study._study_id + status, _, _ = send_request( + app, + f"/api/studies/{study_id}/preference", + "POST", + body=json.dumps({"best_trials": [0, 2], "worst_trials": [1]}), + content_type="application/json", + ) + self.assertEqual(status, 204) + + preferences = study.get_preferences() + preferences.sort(key=lambda x: (x[0].number, x[1].number)) + assert len(preferences) == 2 + better, worse = preferences[0] + assert better.number == 0 + assert worse.number == 1 + better, worse = preferences[1] + assert better.number == 2 + assert worse.number == 1 + def test_create_study(self) -> None: for name, directions, expected_status in [ ("single-objective success", ["minimize"], 201),