Merge pull request #559 from moririn2528/preferential-best-trials

Update JSON APIs to report preferences and get preferential best trials.

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/<study_id:int>/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/<trial_id:int>/tell")
     @json_api_view
     def tell_trial(trial_id: int) -> dict[str, Any]:

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,

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

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

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