diff --git a/optuna_dashboard/_cached_extra_study_property.py b/optuna_dashboard/_cached_extra_study_property.py index 1e9de3fc..4395dad3 100644 --- a/optuna_dashboard/_cached_extra_study_property.py +++ b/optuna_dashboard/_cached_extra_study_property.py @@ -21,7 +21,7 @@ cached_extra_study_property_cache: Dict[int, "_CachedExtraStudyProperty"] = {} def get_cached_extra_study_property( study_id: int, trials: List[FrozenTrial] -) -> Tuple[SearchSpaceListT, SearchSpaceListT, List[str], bool]: +) -> Tuple[SearchSpaceListT, SearchSpaceListT, List[Tuple[str, bool]], bool]: with cached_extra_study_property_cache_lock: cached_extra_study_property = cached_extra_study_property_cache.get(study_id, None) if cached_extra_study_property is None: @@ -43,7 +43,7 @@ class _CachedExtraStudyProperty: # union_user_attrs. self._intersection: Optional[SearchSpaceSetT] = None self._union: SearchSpaceSetT = set() - self._union_user_attrs: Set[str] = set() + self._union_user_attrs: Dict[str, bool] = {} # attr_name: is_sortable (= is_number) self.has_intermediate_values: bool = False @property @@ -61,9 +61,9 @@ class _CachedExtraStudyProperty: return union @property - def union_user_attrs(self) -> List[str]: - union = list(self._union_user_attrs) - union.sort() + def union_user_attrs(self) -> List[Tuple[str, bool]]: + union = [(name, is_sortable) for name, is_sortable in self._union_user_attrs.items()] + sorted(union, key=lambda x: x[0]) return union def update(self, trials: List[FrozenTrial]) -> None: @@ -75,20 +75,27 @@ class _CachedExtraStudyProperty: if not trial.state.is_finished(): next_cursor = trial.number - current_user_attrs = set(n for n in trial.user_attrs.keys()) - self._union_user_attrs = self._union_user_attrs.union(current_user_attrs) - if trial.state == TrialState.FAIL: - continue - - if not self.has_intermediate_values and len(trial.intermediate_values) > 0: - self.has_intermediate_values = True - - current = set([(n, d) for n, d in trial.distributions.items()]) - self._union = self._union.union(current) - - if self._intersection is None: - self._intersection = copy.copy(current) - else: - self._intersection = self._intersection.intersection(current) + self._update_user_attrs(trial) + if trial.state != TrialState.FAIL: + self._update_intermediate_values(trial) + self._update_search_space(trial) self._cursor = next_cursor + + def _update_user_attrs(self, trial: FrozenTrial) -> None: + # TODO(c-bata): Support numpy-specific number types. + current_user_attrs = {k: isinstance(v, (int, float)) for k, v in trial.user_attrs.items()} + self._union_user_attrs.update(current_user_attrs) + + def _update_intermediate_values(self, trial: FrozenTrial) -> None: + if not self.has_intermediate_values and len(trial.intermediate_values) > 0: + self.has_intermediate_values = True + + def _update_search_space(self, trial: FrozenTrial) -> None: + current = set([(n, d) for n, d in trial.distributions.items()]) + self._union = self._union.union(current) + + if self._intersection is None: + self._intersection = copy.copy(current) + else: + self._intersection = self._intersection.intersection(current) diff --git a/optuna_dashboard/_serializer.py b/optuna_dashboard/_serializer.py index 9dca618e..5ed4a027 100644 --- a/optuna_dashboard/_serializer.py +++ b/optuna_dashboard/_serializer.py @@ -29,6 +29,13 @@ Attribute = TypedDict( "value": str, }, ) +AttributeSpec = TypedDict( + "AttributeSpec", + { + "key": str, + "sortable": bool, + }, +) IntermediateValue = TypedDict( "IntermediateValue", { @@ -71,7 +78,7 @@ def serialize_study_detail( trials: List[FrozenTrial], intersection: List[Tuple[str, BaseDistribution]], union: List[Tuple[str, BaseDistribution]], - union_user_attrs: List[str], + union_user_attrs: List[Tuple[str, bool]], has_intermediate_values: bool, ) -> Dict[str, Any]: serialized: Dict[str, Any] = { @@ -84,7 +91,7 @@ def serialize_study_detail( serialized["trials"] = [serialize_frozen_trial(summary._study_id, trial) for trial in trials] serialized["intersection_search_space"] = serialize_search_space(intersection) serialized["union_search_space"] = serialize_search_space(union) - serialized["union_user_attrs"] = union_user_attrs + serialized["union_user_attrs"] = [{"key": a[0], "sortable": a[1]} for a in union_user_attrs] serialized["has_intermediate_values"] = has_intermediate_values serialized["note"] = note.get_note_from_system_attrs(summary.system_attrs) return serialized diff --git a/optuna_dashboard/ts/apiClient.ts b/optuna_dashboard/ts/apiClient.ts index 2f00395b..a539c590 100644 --- a/optuna_dashboard/ts/apiClient.ts +++ b/optuna_dashboard/ts/apiClient.ts @@ -44,7 +44,7 @@ interface StudyDetailResponse { trials: TrialResponse[] intersection_search_space: SearchSpace[] union_search_space: SearchSpace[] - union_user_attrs: string[] + union_user_attrs: AttributeSpec[] has_intermediate_values: boolean note: { version: number diff --git a/optuna_dashboard/ts/components/TrialTable.tsx b/optuna_dashboard/ts/components/TrialTable.tsx index 6763323c..d46c93f1 100644 --- a/optuna_dashboard/ts/components/TrialTable.tsx +++ b/optuna_dashboard/ts/components/TrialTable.tsx @@ -165,21 +165,21 @@ export const TrialTable: FC<{ studyDetail: StudyDetail | null }> = ({ }) } - studyDetail?.union_user_attrs.forEach((attr_name) => { + studyDetail?.union_user_attrs.forEach((attr_spec) => { columns.push({ field: "user_attrs", - label: `User attribute ${attr_name}`, + label: `User attribute ${attr_spec.key}`, toCellValue: (i) => - trials[i].user_attrs.find((attr) => attr.key === attr_name)?.value || - null, - sortable: true, - filterable: true, + trials[i].user_attrs.find((attr) => attr.key === attr_spec.key) + ?.value || null, + sortable: attr_spec.sortable, + filterable: !attr_spec.sortable, less: (firstEl, secondEl): number => { const firstVal = firstEl.user_attrs.find( - (attr) => attr.key === attr_name + (attr) => attr.key === attr_spec.key )?.value const secondVal = secondEl.user_attrs.find( - (attr) => attr.key === attr_name + (attr) => attr.key === attr_spec.key )?.value if (firstVal === secondVal) { diff --git a/optuna_dashboard/ts/types/index.d.ts b/optuna_dashboard/ts/types/index.d.ts index 372f7a5f..6541a93b 100644 --- a/optuna_dashboard/ts/types/index.d.ts +++ b/optuna_dashboard/ts/types/index.d.ts @@ -58,6 +58,11 @@ declare interface Attribute { value: string } +declare interface AttributeSpec { + key: string + sortable: boolean +} + declare interface Note { version: number body: string @@ -95,7 +100,7 @@ declare interface StudyDetail { trials: Trial[] intersection_search_space: SearchSpace[] union_search_space: SearchSpace[] - union_user_attrs: string[] + union_user_attrs: AttributeSpec[] has_intermediate_values: boolean note: Note } diff --git a/typescript_tests/TrialTable.test.tsx b/typescript_tests/TrialTable.test.tsx index 8e87b5d8..da9bdeca 100644 --- a/typescript_tests/TrialTable.test.tsx +++ b/typescript_tests/TrialTable.test.tsx @@ -74,7 +74,10 @@ const studyDetail = { attributes: { low: -3, high: 3 }, }, ], - union_user_attrs: ["foo", "bar"], + union_user_attrs: [ + { key: "foo", sortable: false }, + { key: "bar", sortable: false }, + ], has_intermediate_values: false, note: { version: 0,