def get_all_study_summaries(self) -> List[StudySummary]:
session = self.scoped_session()
summarized_trial = (session.query(
models.TrialModel.study_id,
functions.min(
models.TrialModel.datetime_start).label("datetime_start"),
functions.count(models.TrialModel.trial_id).label("n_trial"),
).group_by(models.TrialModel.study_id).with_labels().subquery())
study_summary_stmt = session.query(
models.StudyModel.study_id,
models.StudyModel.study_name,
models.StudyModel.direction,
summarized_trial.c.datetime_start,
functions.coalesce(summarized_trial.c.n_trial, 0).label("n_trial"),
).select_from(orm.outerjoin(models.StudyModel, summarized_trial))
study_summary = study_summary_stmt.all()
study_summaries = []
for study in study_summary:
best_trial = None # type: Optional[models.TrialModel]
try:
if study.direction == StudyDirection.MAXIMIZE:
best_trial = models.TrialModel.find_max_value_trial(
study.study_id, session)
else:
best_trial = models.TrialModel.find_min_value_trial(
study.study_id, session)
except ValueError:
best_trial_frozen = None # type: Optional[FrozenTrial]
if best_trial:
params = (session.query(
models.TrialParamModel.param_name,
models.TrialParamModel.param_value,
models.TrialParamModel.distribution_json,
).filter(models.TrialParamModel.trial_id ==
best_trial.trial_id).all())
param_dict = {}
param_distributions = {}
for param in params:
distribution = distributions.json_to_distribution(
param.distribution_json)
param_dict[
param.param_name] = distribution.to_external_repr(
param.param_value)
param_distributions[param.param_name] = distribution
user_attrs = session.query(
models.TrialUserAttributeModel).filter(
models.TrialUserAttributeModel.trial_id ==
best_trial.trial_id)
system_attrs = session.query(
models.TrialSystemAttributeModel).filter(
models.TrialSystemAttributeModel.trial_id ==
best_trial.trial_id)
intermediate = session.query(models.TrialValueModel).filter(
models.TrialValueModel.trial_id == best_trial.trial_id)
best_trial_frozen = FrozenTrial(
best_trial.number,
TrialState.COMPLETE,
best_trial.value,
best_trial.datetime_start,
best_trial.datetime_complete,
param_dict,
param_distributions,
{i.key: json.loads(i.value_json)
for i in user_attrs},
{i.key: json.loads(i.value_json)
for i in system_attrs},
{value.step: value.value
for value in intermediate},
best_trial.trial_id,
)
user_attrs = session.query(models.StudyUserAttributeModel).filter(
models.StudyUserAttributeModel.study_id == study.study_id)
system_attrs = session.query(
models.StudySystemAttributeModel).filter(
models.StudySystemAttributeModel.study_id ==
study.study_id)
study_summaries.append(
StudySummary(
study_name=study.study_name,
direction=study.direction,
best_trial=best_trial_frozen,
user_attrs={
i.key: json.loads(i.value_json)
for i in user_attrs
},
system_attrs={
i.key: json.loads(i.value_json)
for i in system_attrs
},
n_trials=study.n_trial,
datetime_start=study.datetime_start,
study_id=study.study_id,
))
# Terminate transaction explicitly to avoid connection timeout during transaction.
self._commit(session)
return study_summaries
EXAMPLE_DISTRIBUTIONS = {
"x": UniformDistribution(low=1.0, high=2.0),
"y": CategoricalDistribution(choices=("Otemachi", "Tokyo", "Ginza")),
} # type: Dict[str, BaseDistribution]
EXAMPLE_TRIALS = [
FrozenTrial(
number=0, # dummy
value=1.0,
state=TrialState.COMPLETE,
user_attrs={},
system_attrs={},
params={
"x": 0.5,
"y": "Ginza"
},
distributions=EXAMPLE_DISTRIBUTIONS,
intermediate_values={
0: 2.0,
1: 3.0
},
datetime_start=None, # dummy
datetime_complete=None, # dummy
trial_id=-1, # dummy id
),
FrozenTrial(
number=0, # dummy
value=2.0,
state=TrialState.RUNNING,
user_attrs={
"tags": ["video", "classification"],
def test_update_trial_second_write() -> None:
storage = create_test_storage()
study_id = storage.create_new_study()
template = FrozenTrial(
number=1,
state=TrialState.RUNNING,
value=0.1,
datetime_start=None,
datetime_complete=None,
params={
"paramA": 0.1,
"paramB": 1.1
},
distributions={
"paramA": UniformDistribution(0, 1),
"paramB": UniformDistribution(0, 2)
},
user_attrs={
"userA": 2,
"userB": 3
},
system_attrs={
"sysA": 4,
"sysB": 5
},
intermediate_values={
3: 1.2,
5: 9.2
},
trial_id=1,
)
trial_id = storage.create_new_trial(study_id, template)
trial_before_update = storage.get_trial(trial_id)
storage._update_trial(
trial_id,
state=None,
value=1.1,
intermediate_values={
3: 2.3,
7: 3.3
},
params={
"paramA": 0.2,
"paramC": 2.3
},
distributions_={
"paramA": UniformDistribution(0, 1),
"paramC": UniformDistribution(0, 4)
},
user_attrs={
"userA": 1,
"userC": "attr"
},
system_attrs={
"sysA": 6,
"sysC": 8
},
)
trial_after_update = storage.get_trial(trial_id)
expected_attrs = {
"_trial_id": trial_before_update._trial_id,
"number": trial_before_update.number,
"state": TrialState.RUNNING,
"value": 1.1,
"params": {
"paramA": 0.2,
"paramB": 1.1,
"paramC": 2.3
},
"intermediate_values": {
3: 2.3,
5: 9.2,
7: 3.3
},
"_distributions": {
"paramA": UniformDistribution(0, 1),
"paramB": UniformDistribution(0, 2),
"paramC": UniformDistribution(0, 4),
},
"user_attrs": {
"userA": 1,
"userB": 3,
"userC": "attr"
},
"system_attrs": {
"sysA": 6,
"sysB": 5,
"sysC": 8
},
}
for key, value in expected_attrs.items():
assert getattr(trial_after_update, key) == value
def _create_new_trial(
self, study_id: int, template_trial: Optional[FrozenTrial] = None
) -> FrozenTrial:
"""Create a new trial and returns its trial_id and a :class:`~optuna.trial.FrozenTrial`.
Args:
study_id:
Study id.
template_trial:
A :class:`~optuna.trial.FrozenTrial` with default values for trial attributes.
Returns:
A :class:`~optuna.trial.FrozenTrial` instance.
"""
# Retry a couple of times. Deadlocks may occur in distributed environments.
n_retries = 0
while True:
session = self.scoped_session()
try:
# Ensure that that study exists.
#
# Locking within a study is necessary since the creation of a trial is not an
# atomic operation. More precisely, the trial number computed in
# `_get_prepared_new_trial` is prone to race conditions without this lock.
models.StudyModel.find_or_raise_by_id(study_id, session, for_update=True)
trial = self._get_prepared_new_trial(study_id, template_trial, session)
self._commit(session)
break # Successfully created trial.
except OperationalError:
session.rollback()
if n_retries > 2:
raise
n_retries += 1
if template_trial:
frozen = copy.deepcopy(template_trial)
frozen.number = trial.number
frozen.datetime_start = trial.datetime_start
frozen._trial_id = trial.trial_id
else:
frozen = FrozenTrial(
number=trial.number,
state=trial.state,
value=None,
datetime_start=trial.datetime_start,
datetime_complete=None,
params={},
distributions={},
user_attrs={},
system_attrs={},
intermediate_values={},
trial_id=trial.trial_id,
)
return frozen
def _create_new_trial(
self,
study_id: int,
template_trial: Optional[FrozenTrial] = None) -> FrozenTrial:
"""Create a new trial and returns its trial_id and a :class:`~optuna.trial.FrozenTrial`.
Args:
study_id:
Study id.
template_trial:
A :class:`~optuna.trial.FrozenTrial` with default values for trial attributes.
Returns:
A :class:`~optuna.trial.FrozenTrial` instance.
"""
session = self.scoped_session()
# Ensure that that study exists.
models.StudyModel.find_or_raise_by_id(study_id, session)
if template_trial is None:
trial = models.TrialModel(study_id=study_id,
number=None,
state=TrialState.RUNNING)
else:
# Because only `RUNNING` trials can be updated,
# we temporarily set the state of the new trial to `RUNNING`.
# After all fields of the trial have been updated,
# the state is set to `template_trial.state`.
temp_state = TrialState.RUNNING
trial = models.TrialModel(
study_id=study_id,
number=None,
state=temp_state,
value=template_trial.value,
datetime_start=template_trial.datetime_start,
datetime_complete=template_trial.datetime_complete,
)
session.add(trial)
# Flush the session cache to reflect the above addition operation to
# the current RDB transaction.
#
# Without flushing, the following operations (e.g, `_set_trial_param_without_commit`)
# will fail because the target trial doesn't exist in the storage yet.
session.flush()
if template_trial is not None:
for param_name, param_value in template_trial.params.items():
distribution = template_trial.distributions[param_name]
param_value_in_internal_repr = distribution.to_internal_repr(
param_value)
self._set_trial_param_without_commit(
session, trial.trial_id, param_name,
param_value_in_internal_repr, distribution)
for key, value in template_trial.user_attrs.items():
self._set_trial_user_attr_without_commit(
session, trial.trial_id, key, value)
for key, value in template_trial.system_attrs.items():
self._set_trial_system_attr_without_commit(
session, trial.trial_id, key, value)
for step, intermediate_value in template_trial.intermediate_values.items(
):
self._set_trial_intermediate_value_without_commit(
session, trial.trial_id, step, intermediate_value)
trial.state = template_trial.state
trial.number = trial.count_past_trials(session)
session.add(trial)
self._commit(session)
if template_trial:
frozen = copy.deepcopy(template_trial)
frozen.number = trial.number
frozen.datetime_start = trial.datetime_start
frozen._trial_id = trial.trial_id
else:
frozen = FrozenTrial(
number=trial.number,
state=trial.state,
value=None,
datetime_start=trial.datetime_start,
datetime_complete=None,
params={},
distributions={},
user_attrs={},
system_attrs={},
intermediate_values={},
trial_id=trial.trial_id,
)
return frozen
def test_make_hovertext() -> None:
trial_no_user_attrs = FrozenTrial(
number=0,
trial_id=0,
state=TrialState.COMPLETE,
value=0.2,
datetime_start=datetime.datetime.now(),
datetime_complete=datetime.datetime.now(),
params={"x": 10},
distributions={"x": FloatDistribution(5, 12)},
user_attrs={},
system_attrs={},
intermediate_values={},
)
assert (_make_hovertext(trial_no_user_attrs) == dedent("""
{
"number": 0,
"values": [
0.2
],
"params": {
"x": 10
}
}
""").strip().replace("\n", "<br>"))
trial_user_attrs_valid_json = FrozenTrial(
number=0,
trial_id=0,
state=TrialState.COMPLETE,
value=0.2,
datetime_start=datetime.datetime.now(),
datetime_complete=datetime.datetime.now(),
params={"x": 10},
distributions={"x": FloatDistribution(5, 12)},
user_attrs={
"a": 42,
"b": 3.14
},
system_attrs={},
intermediate_values={},
)
assert (_make_hovertext(trial_user_attrs_valid_json) == dedent("""
{
"number": 0,
"values": [
0.2
],
"params": {
"x": 10
},
"user_attrs": {
"a": 42,
"b": 3.14
}
}
""").strip().replace("\n", "<br>"))
trial_user_attrs_invalid_json = FrozenTrial(
number=0,
trial_id=0,
state=TrialState.COMPLETE,
value=0.2,
datetime_start=datetime.datetime.now(),
datetime_complete=datetime.datetime.now(),
params={"x": 10},
distributions={"x": FloatDistribution(5, 12)},
user_attrs={
"a": 42,
"b": 3.14,
"c": np.zeros(1),
"d": np.nan
},
system_attrs={},
intermediate_values={},
)
assert (_make_hovertext(trial_user_attrs_invalid_json) == dedent("""
{
"number": 0,
"values": [
0.2
],
"params": {
"x": 10
},
"user_attrs": {
"a": 42,
"b": 3.14,
"c": "[0.]",
"d": NaN
}
}
""").strip().replace("\n", "<br>"))
def test_create_new_trial_with_template_trial(storage_mode: str) -> None:
start_time = datetime.now()
complete_time = datetime.now()
template_trial = FrozenTrial(
state=TrialState.COMPLETE,
value=10000,
datetime_start=start_time,
datetime_complete=complete_time,
params={"x": 0.5},
distributions={"x": UniformDistribution(0, 1)},
user_attrs={"foo": "bar"},
system_attrs={"baz": 123},
intermediate_values={
1: 10,
2: 100,
3: 1000
},
number=55, # This entry is ignored.
trial_id=-1, # dummy value (unused).
)
def _check_trials(trials: List[FrozenTrial], idx: int,
trial_id: int) -> None:
assert len(trials) == idx + 1
assert len({t._trial_id for t in trials}) == idx + 1
assert trial_id in {t._trial_id for t in trials}
assert {t.number for t in trials} == set(range(idx + 1))
assert all(t.state == template_trial.state for t in trials)
assert all(t.params == template_trial.params for t in trials)
assert all(t.distributions == template_trial.distributions
for t in trials)
assert all(t.intermediate_values == template_trial.intermediate_values
for t in trials)
assert all(t.user_attrs == template_trial.user_attrs for t in trials)
assert all(t.system_attrs == template_trial.system_attrs
for t in trials)
assert all(t.datetime_start == template_trial.datetime_start
for t in trials)
assert all(t.datetime_complete == template_trial.datetime_complete
for t in trials)
assert all(t.value == template_trial.value for t in trials)
with StorageSupplier(storage_mode) as storage:
study_id = storage.create_new_study()
n_trial_in_study = 3
for i in range(n_trial_in_study):
trial_id = storage.create_new_trial(study_id,
template_trial=template_trial)
trials = storage.get_all_trials(study_id)
_check_trials(trials, i, trial_id)
# Create trial in non-existent study.
with pytest.raises(KeyError):
storage.create_new_trial(study_id + 1)
study_id2 = storage.create_new_study()
for i in range(n_trial_in_study):
storage.create_new_trial(study_id2, template_trial=template_trial)
trials = storage.get_all_trials(study_id2)
assert {t.number for t in trials} == set(range(i + 1))
trials = storage.get_all_trials(study_id) + storage.get_all_trials(
study_id2)
# Check trial_ids are unique across studies.
assert len({t._trial_id for t in trials}) == 2 * n_trial_in_study
def _merge_trials_orm(
self,
trials, # type: List[models.TrialModel]
trial_params, # type: List[models.TrialParamModel]
trial_intermediate_values, # type: List[models.TrialValueModel]
trial_user_attrs, # type: List[models.TrialUserAttributeModel]
trial_system_attrs, # type: List[models.TrialSystemAttributeModel]
):
# type: (...) -> List[FrozenTrial]
id_to_trial = {}
for trial in trials:
id_to_trial[trial.trial_id] = trial
id_to_params = defaultdict(
list) # type: Dict[int, List[models.TrialParamModel]]
for param in trial_params:
id_to_params[param.trial_id].append(param)
id_to_values = defaultdict(
list) # type: Dict[int, List[models.TrialValueModel]]
for value in trial_intermediate_values:
id_to_values[value.trial_id].append(value)
id_to_user_attrs = defaultdict(
list) # type: Dict[int, List[models.TrialUserAttributeModel]]
for user_attr in trial_user_attrs:
id_to_user_attrs[user_attr.trial_id].append(user_attr)
id_to_system_attrs = defaultdict(
list) # type: Dict[int, List[models.TrialSystemAttributeModel]]
for system_attr in trial_system_attrs:
id_to_system_attrs[system_attr.trial_id].append(system_attr)
result = []
for trial_id, trial in id_to_trial.items():
params = {}
param_distributions = {}
for param in id_to_params[trial_id]:
distribution = distributions.json_to_distribution(
param.distribution_json)
params[param.param_name] = distribution.to_external_repr(
param.param_value)
param_distributions[param.param_name] = distribution
intermediate_values = {}
for value in id_to_values[trial_id]:
intermediate_values[value.step] = value.value
user_attrs = {}
for user_attr in id_to_user_attrs[trial_id]:
user_attrs[user_attr.key] = json.loads(user_attr.value_json)
system_attrs = {}
for system_attr in id_to_system_attrs[trial_id]:
system_attrs[system_attr.key] = json.loads(
system_attr.value_json)
result.append(
FrozenTrial(
number=trial.number,
state=trial.state,
params=params,
distributions=param_distributions,
user_attrs=user_attrs,
system_attrs=system_attrs,
value=trial.value,
intermediate_values=intermediate_values,
datetime_start=trial.datetime_start,
datetime_complete=trial.datetime_complete,
trial_id=trial_id,
))
return result
