Files
optuna-dashboard/optuna_dashboard/_serializer.py
T
2024-01-24 10:06:40 +09:00

386 lines
14 KiB
Python

from __future__ import annotations
from datetime import datetime
import json
import numbers
from typing import Any
from typing import TYPE_CHECKING
from typing import Union
import numpy as np
from optuna.distributions import BaseDistribution
from optuna.distributions import CategoricalDistribution
from optuna.study import StudySummary
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_COMPONENT
from ._preferential_history import _SYSTEM_ATTR_PREFIX_HISTORY
from .artifact._backend import list_study_artifacts
from .artifact._backend import list_trial_artifacts
from .preferential._study import _SYSTEM_ATTR_PREFERENTIAL_STUDY
from .preferential._system_attrs import get_preferences
from .preferential._system_attrs import is_preference_removed
if TYPE_CHECKING:
from typing import Literal
from typing import TypedDict
from ._preferential_history import History
from ._preferential_history import SerializedHistory
Attribute = TypedDict(
"Attribute",
{
"key": str,
"value": str,
},
)
AttributeSpec = TypedDict(
"AttributeSpec",
{
"key": str,
"sortable": bool,
},
)
IntermediateValue = TypedDict(
"IntermediateValue",
{
"step": int,
"value": Union[float, Literal["inf", "-inf", "nan"]],
},
)
FloatDistributionJSON = TypedDict(
"FloatDistributionJSON",
{
"type": Literal["FloatDistribution"],
"low": float,
"high": float,
"step": float,
"log": bool,
},
)
IntDistributionJSON = TypedDict(
"IntDistributionJSON",
{
"type": Literal["IntDistribution"],
"low": int,
"high": int,
"step": int,
"log": bool,
},
)
CategoricalDistributionChoiceJSON = TypedDict(
"CategoricalDistributionChoiceJSON",
{
"pytype": str,
"value": str,
},
)
CategoricalDistributionJSON = TypedDict(
"CategoricalDistributionJSON",
{
"type": Literal["CategoricalDistribution"],
"choices": list[CategoricalDistributionChoiceJSON],
},
)
DistributionJSON = Union[
FloatDistributionJSON, IntDistributionJSON, CategoricalDistributionJSON
]
MAX_ATTR_LENGTH = 1024
CONSTRAINTS_KEY = "constraints"
def serialize_attrs(attrs: dict[str, Any]) -> list[Attribute]:
serialized: list[Attribute] = []
for k, v in attrs.items():
value: str
if isinstance(v, bytes):
value = "<binary object>"
elif isinstance(v, str):
value = v
elif isinstance(v, numbers.Real):
value = str(v)
else:
value = json.dumps(v)
value = value[:MAX_ATTR_LENGTH] if len(value) > MAX_ATTR_LENGTH else value
serialized.append({"key": k, "value": value})
return serialized
def serialize_study_summary(summary: StudySummary) -> dict[str, Any]:
serialized = {
"study_id": summary._study_id,
"study_name": summary.study_name,
"directions": [d.name.lower() for d in summary.directions],
"user_attrs": serialize_attrs(summary.user_attrs),
"is_preferential": getattr(summary, "_system_attrs", {}).get(
_SYSTEM_ATTR_PREFERENTIAL_STUDY, False
),
}
if summary.datetime_start is not None:
serialized["datetime_start"] = summary.datetime_start.isoformat()
return serialized
def serialize_study_detail(
summary: StudySummary,
best_trials: list[FrozenTrial],
trials: list[FrozenTrial],
intersection: list[tuple[str, BaseDistribution]],
union: list[tuple[str, BaseDistribution]],
union_user_attrs: list[tuple[str, bool]],
has_intermediate_values: bool,
plotly_graph_objects: dict[str, str],
skipped_trial_numbers: list[int],
) -> dict[str, Any]:
serialized: dict[str, Any] = {
"name": summary.study_name,
"directions": [d.name.lower() for d in summary.directions],
"user_attrs": serialize_attrs(summary.user_attrs),
}
system_attrs = getattr(summary, "system_attrs", {})
serialized["artifacts"] = list_study_artifacts(system_attrs)
if summary.datetime_start is not None:
serialized["datetime_start"] = summary.datetime_start.isoformat()
serialized["trials"] = [
serialize_frozen_trial(summary._study_id, trial, system_attrs) for trial in trials
]
serialized["best_trials"] = [
serialize_frozen_trial(summary._study_id, trial, system_attrs) for trial in best_trials
]
serialized["intersection_search_space"] = serialize_search_space(intersection)
serialized["union_search_space"] = serialize_search_space(union)
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(system_attrs, None)
serialized["is_preferential"] = system_attrs.get(_SYSTEM_ATTR_PREFERENTIAL_STUDY, False)
objective_names = get_objective_names(system_attrs)
if objective_names:
serialized["objective_names"] = objective_names
form_widgets = get_form_widgets_json(system_attrs)
if form_widgets:
serialized["form_widgets"] = form_widgets
serialized["feedback_component_type"] = system_attrs.get(
_SYSTEM_ATTR_FEEDBACK_COMPONENT,
{
"output_type": "note",
},
)
if serialized["is_preferential"]:
serialized["preference_history"] = serialize_preference_history(system_attrs)
serialized["preferences"] = get_preferences(system_attrs)
serialized["skipped_trial_numbers"] = skipped_trial_numbers
serialized["plotly_graph_objects"] = [
{"id": id_, "graph_object": graph_object}
for id_, graph_object in plotly_graph_objects.items()
]
return serialized
def serialize_preference_history(
system_attrs: dict[str, Any],
) -> list[SerializedHistory]:
histories: list[SerializedHistory] = []
for k, v in system_attrs.items():
if not k.startswith(_SYSTEM_ATTR_PREFIX_HISTORY):
continue
choice: dict[str, Any] = json.loads(v)
if choice["mode"] == "ChooseWorst":
history: History = {
"mode": "ChooseWorst",
"id": choice["id"],
"timestamp": choice["timestamp"],
"candidates": choice["candidates"],
"clicked": choice["clicked"],
"preferences": choice["preferences"],
}
histories.append(
{
"history": history,
"is_removed": is_preference_removed(system_attrs, choice["id"]),
}
)
histories.sort(key=lambda c: datetime.fromisoformat(c["history"]["timestamp"]))
return histories
def serialize_frozen_trial(
study_id: int, trial: FrozenTrial, study_system_attrs: dict[str, Any]
) -> dict[str, Any]:
params = []
for param_name, param_external_value in trial.params.items():
distribution = trial.distributions.get(param_name)
if distribution is None:
continue
params.append(
{
"name": param_name,
"param_internal_value": distribution.to_internal_repr(param_external_value),
"param_external_value": str(param_external_value),
"param_external_pytyp": str(type(param_external_value)),
"distribution": serialize_distribution(distribution),
}
)
trial_system_attrs: dict[str, Any] = getattr(trial, "_system_attrs", {})
fixed_params = trial_system_attrs.get("fixed_params", {})
serialized = {
"trial_id": trial._trial_id,
"study_id": study_id,
"number": trial.number,
"state": trial.state.name.capitalize(),
"params": params,
"fixed_params": [
{"name": param_name, "param_external_value": str(fixed_params.get(param_name, None))}
for param_name in fixed_params
],
"user_attrs": serialize_attrs(trial.user_attrs),
"note": note.get_note_from_system_attrs(study_system_attrs, trial._trial_id),
"artifacts": list_trial_artifacts(study_system_attrs, trial_system_attrs, trial),
"constraints": trial_system_attrs.get(CONSTRAINTS_KEY, []),
}
serialized_intermediate_values: list[IntermediateValue] = []
for step, value in trial.intermediate_values.items():
serialized_value: Union[float, Literal["nan", "inf", "-inf"]]
if np.isnan(value):
serialized_value = "nan"
elif np.isposinf(value):
serialized_value = "inf"
elif np.isneginf(value):
serialized_value = "-inf"
else:
assert np.isfinite(value)
serialized_value = value
serialized_intermediate_values.append({"step": step, "value": serialized_value})
serialized["intermediate_values"] = sorted(
serialized_intermediate_values, key=lambda v: v["step"]
)
if trial.values is not None:
serialized_values: list[Union[float, Literal["inf", "-inf"]]] = []
for v in trial.values:
assert not np.isnan(v), "Should not detect nan value"
if np.isposinf(v):
serialized_values.append("inf")
elif np.isneginf(v):
serialized_values.append("-inf")
else:
serialized_values.append(v)
serialized["values"] = serialized_values
if trial.datetime_start is not None:
serialized["datetime_start"] = trial.datetime_start.isoformat()
if trial.datetime_complete is not None:
serialized["datetime_complete"] = trial.datetime_complete.isoformat()
return serialized
def serialize_distribution(distribution: BaseDistribution) -> DistributionJSON:
if distribution.__class__.__name__ == "FloatDistribution":
# Added from Optuna v3.0
float_distribution: FloatDistributionJSON = {
"type": "FloatDistribution",
"low": getattr(distribution, "low"),
"high": getattr(distribution, "high"),
"step": getattr(distribution, "step"),
"log": getattr(distribution, "log"),
}
return float_distribution
if distribution.__class__.__name__ == "UniformDistribution":
# Deprecated from Optuna v3.0
uniform: FloatDistributionJSON = {
"type": "FloatDistribution",
"low": getattr(distribution, "low"),
"high": getattr(distribution, "high"),
"step": 0,
"log": False,
}
return uniform
if distribution.__class__.__name__ == "LogUniformDistribution":
# Deprecated from Optuna v3.0
log_uniform: FloatDistributionJSON = {
"type": "FloatDistribution",
"low": getattr(distribution, "low"),
"high": getattr(distribution, "high"),
"step": 0,
"log": True,
}
return log_uniform
if distribution.__class__.__name__ == "DiscreteUniformDistribution":
# Deprecated from Optuna v3.0
discrete_uniform: FloatDistributionJSON = {
"type": "FloatDistribution",
"low": getattr(distribution, "low"),
"high": getattr(distribution, "high"),
"step": getattr(distribution, "q"),
"log": False,
}
return discrete_uniform
if distribution.__class__.__name__ == "IntDistribution":
# Added from Optuna v3.0
int_distribution: IntDistributionJSON = {
"type": "IntDistribution",
"low": getattr(distribution, "low"),
"high": getattr(distribution, "high"),
"step": getattr(distribution, "step"),
"log": getattr(distribution, "log"),
}
return int_distribution
if distribution.__class__.__name__ == "IntUniformDistribution":
# Deprecated from Optuna v3.0
int_uniform: IntDistributionJSON = {
"type": "IntDistribution",
"low": getattr(distribution, "low"),
"high": getattr(distribution, "high"),
"step": getattr(distribution, "step"),
"log": False,
}
return int_uniform
if distribution.__class__.__name__ == "IntLogUniformDistribution":
# Deprecated from Optuna v3.0
int_log_uniform: IntDistributionJSON = {
"type": "IntDistribution",
"low": getattr(distribution, "low"),
"high": getattr(distribution, "high"),
"step": getattr(distribution, "step"),
"log": True,
}
return int_log_uniform
if isinstance(distribution, CategoricalDistribution):
categorical: CategoricalDistributionJSON = {
"type": "CategoricalDistribution",
"choices": [
{"pytype": str(type(choice)), "value": str(choice)}
for choice in distribution.choices
],
}
return categorical
raise ValueError(f"Unexpected distribution {str(distribution)}")
def serialize_search_space(
search_space: list[tuple[str, BaseDistribution]]
) -> list[dict[str, Any]]:
serialized = []
for param_name, distribution in search_space:
serialized.append(
{
"name": param_name,
"distribution": serialize_distribution(distribution),
}
)
return serialized