diff --git a/optuna_dashboard/_app.py b/optuna_dashboard/_app.py index c32c2061..fad4b74c 100644 --- a/optuna_dashboard/_app.py +++ b/optuna_dashboard/_app.py @@ -28,6 +28,7 @@ from ._cached_extra_study_property import get_cached_extra_study_property from ._custom_plot_data import get_plotly_graph_objects from ._importance import get_param_importance_from_trials_cache from ._pareto_front import get_pareto_front_trials +from ._preference_setting import _register_output_component from ._preferential_history import NewHistory from ._preferential_history import report_history from ._rdb_migration import register_rdb_migration_route @@ -306,6 +307,28 @@ def create_app( response.status = 204 return {} + @app.post("/api/studies//component") + @json_api_view + def post_component(study_id: int) -> dict[str, Any]: + try: + component_type = request.json.get("component_type", "") + artifact_key = request.json.get("artifact_key", None) + except ValueError: + response.status = 400 + return {"reason": "invalid request."} + if component_type not in ["Note", "Artifact"]: + response.status = 400 + return {"reason": "component_type must be either 'Note' or 'Artifact'."} + + _register_output_component( + study_id=study_id, + storage=storage, + component_type=component_type, + artifact_key=artifact_key, + ) + 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/_preference_setting.py b/optuna_dashboard/_preference_setting.py new file mode 100644 index 00000000..93b1a34c --- /dev/null +++ b/optuna_dashboard/_preference_setting.py @@ -0,0 +1,60 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +from optuna.storages import BaseStorage + +from .preferential._study import PreferentialStudy + + +if TYPE_CHECKING: + from typing import Literal + + OUTPUT_COMPONENT_TYPE = Literal["Note", "Artifact"] + +_SYSTEM_ATTR_FEEDBACK_COMPONENT_TYPE = "preference:component_type" +_SYSTEM_ATTR_FEEDBACK_ARTIFACT_KEY = "preference:component_artifact_key" + + +def _register_output_component( + study_id: int, + storage: BaseStorage, + component_type: OUTPUT_COMPONENT_TYPE, + artifact_key: str | None = None, +) -> None: + storage.set_study_system_attr( + study_id=study_id, + key=_SYSTEM_ATTR_FEEDBACK_COMPONENT_TYPE, + value=component_type, + ) + if artifact_key is not None: + storage.set_study_system_attr( + study_id=study_id, + key=_SYSTEM_ATTR_FEEDBACK_ARTIFACT_KEY, + value=artifact_key, + ) + + +def register_output_component( + study: PreferentialStudy, + component_type: OUTPUT_COMPONENT_TYPE, + artifact_key: str = "", +) -> None: + """Register output component to the study. + + Args: + study: + The study to register the output component. + component_type: + The type of the output component. + artifact_key: + When the component_type is "Artifact", + this argument is used as the attribute key of the artifact. + Each trial displays the artifact whose id is the value of the attribute. + """ + _register_output_component( + study_id=study._study._study_id, + storage=study._study._storage, + component_type=component_type, + artifact_key=artifact_key, + ) diff --git a/optuna_dashboard/_serializer.py b/optuna_dashboard/_serializer.py index e3f77649..511e7365 100644 --- a/optuna_dashboard/_serializer.py +++ b/optuna_dashboard/_serializer.py @@ -15,6 +15,8 @@ from optuna.trial import FrozenTrial from . import _note as note from ._form_widget import get_form_widgets_json from ._named_objectives import get_objective_names +from ._preference_setting import _SYSTEM_ATTR_FEEDBACK_ARTIFACT_KEY +from ._preference_setting import _SYSTEM_ATTR_FEEDBACK_COMPONENT_TYPE from ._preferential_history import _SYSTEM_ATTR_PREFIX_HISTORY from .artifact._backend import list_trial_artifacts from .preferential._study import _SYSTEM_ATTR_PREFERENTIAL_STUDY @@ -162,6 +164,10 @@ def serialize_study_detail( form_widgets = get_form_widgets_json(system_attrs) if form_widgets: serialized["form_widgets"] = form_widgets + if _SYSTEM_ATTR_FEEDBACK_COMPONENT_TYPE in system_attrs: + serialized["feedback_component_type"] = system_attrs[_SYSTEM_ATTR_FEEDBACK_COMPONENT_TYPE] + if _SYSTEM_ATTR_FEEDBACK_ARTIFACT_KEY in system_attrs: + serialized["feedback_artifact_key"] = system_attrs[_SYSTEM_ATTR_FEEDBACK_ARTIFACT_KEY] if serialized["is_preferential"]: serialized["preference_history"] = serialize_preference_history(system_attrs) serialized["preferences"] = get_preferences(system_attrs) diff --git a/python_tests/test_api.py b/python_tests/test_api.py index c551e3f2..e22d104c 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._preference_setting import register_output_component from optuna_dashboard.preferential import create_study from .wsgi_client import send_request @@ -179,6 +180,66 @@ class APITestCase(TestCase): ) self.assertEqual(status, 400) + def test_change_component(self) -> None: + storage = optuna.storages.InMemoryStorage() + study = create_study(storage=storage, n_generate=3) + register_output_component(study, "Note") + for _ in range(3): + study.ask() + + app = create_app(storage) + study_id = study._study._study_id + status, _, _ = send_request( + app, + f"/api/studies/{study_id}/component", + "POST", + body=json.dumps({"component_type": "Artifact", "artifact_key": "image"}), + content_type="application/json", + ) + self.assertEqual(status, 204) + + status, _, body = send_request( + app, + f"/api/studies/{study_id}", + "GET", + content_type="application/json", + ) + self.assertEqual(status, 200) + + study_detail = json.loads(body) + assert study_detail["feedback_component_type"] == "Artifact" + assert study_detail["feedback_artifact_key"] == "image" + + def test_change_component_type_only(self) -> None: + storage = optuna.storages.InMemoryStorage() + study = create_study(storage=storage, n_generate=3) + register_output_component(study, "Artifact", "audio") + for _ in range(3): + study.ask() + + app = create_app(storage) + study_id = study._study._study_id + status, _, _ = send_request( + app, + f"/api/studies/{study_id}/component", + "POST", + body=json.dumps({"component_type": "Note"}), + content_type="application/json", + ) + self.assertEqual(status, 204) + + status, _, body = send_request( + app, + f"/api/studies/{study_id}", + "GET", + content_type="application/json", + ) + self.assertEqual(status, 200) + + study_detail = json.loads(body) + assert study_detail["feedback_component_type"] == "Note" + assert study_detail["feedback_artifact_key"] == "audio" + def test_skip_trial(self) -> None: storage = optuna.storages.InMemoryStorage() study = create_study(n_generate=4, storage=storage) diff --git a/python_tests/test_preference_setting.py b/python_tests/test_preference_setting.py new file mode 100644 index 00000000..033a9092 --- /dev/null +++ b/python_tests/test_preference_setting.py @@ -0,0 +1,18 @@ +from __future__ import annotations + +from unittest import TestCase + +import optuna +from optuna_dashboard._preference_setting import _SYSTEM_ATTR_FEEDBACK_ARTIFACT_KEY +from optuna_dashboard._preference_setting import _SYSTEM_ATTR_FEEDBACK_COMPONENT_TYPE +from optuna_dashboard._preference_setting import register_output_component +from optuna_dashboard.preferential._study import PreferentialStudy + + +class FeedbackSettingTestCase(TestCase): + def test_widget_to_dict_from_dict(self) -> None: + study = PreferentialStudy(optuna.create_study()) + register_output_component(study, "Artifact", "image_key") + system_attrs = study._study.system_attrs + assert system_attrs.get(_SYSTEM_ATTR_FEEDBACK_COMPONENT_TYPE, "") == "Artifact" + assert system_attrs.get(_SYSTEM_ATTR_FEEDBACK_ARTIFACT_KEY, "") == "image_key"