def _record_heartbeat(trial_id: int, storage: storages.BaseStorage, stop_event: Event) -> None:
heartbeat_interval = storage.get_heartbeat_interval()
assert heartbeat_interval is not None
while True:
storage.record_heartbeat(trial_id)
if stop_event.wait(timeout=heartbeat_interval):
return
def _setup_studies(
storage: BaseStorage,
n_study: int,
n_trial: int,
seed: int,
direction: Optional[StudyDirection] = None,
) -> Tuple[Dict[int, FrozenStudy], Dict[int, Dict[int, FrozenTrial]]]:
generator = random.Random(seed)
study_id_to_frozen_study: Dict[int, FrozenStudy] = {}
study_id_to_trials: Dict[int, Dict[int, FrozenTrial]] = {}
for i in range(n_study):
study_name = "test-study-name-{}".format(i)
study_id = storage.create_new_study(study_name=study_name)
if direction is None:
direction = generator.choice(
[StudyDirection.MINIMIZE, StudyDirection.MAXIMIZE])
storage.set_study_directions(study_id, (direction, ))
storage.set_study_user_attr(study_id, "u", i)
storage.set_study_system_attr(study_id, "s", i)
trials = {}
for j in range(n_trial):
trial = _generate_trial(generator)
trial.number = j
trial._trial_id = storage.create_new_trial(study_id, trial)
trials[trial._trial_id] = trial
study_id_to_trials[study_id] = trials
study_id_to_frozen_study[study_id] = FrozenStudy(
study_name=study_name,
direction=direction,
user_attrs={"u": i},
system_attrs={"s": i},
study_id=study_id,
)
return study_id_to_frozen_study, study_id_to_trials
def _setup_studies(
storage: BaseStorage,
n_study: int,
n_trial: int,
seed: int,
direction: Optional[StudyDirection] = None,
) -> Tuple[Dict[int, StudySummary], Dict[int, Dict[int, FrozenTrial]]]:
generator = random.Random(seed)
study_id_to_summary = {} # type: Dict[int, StudySummary]
study_id_to_trials = {} # type: Dict[int, Dict[int, FrozenTrial]]
for i in range(n_study):
study_name = "test-study-name-{}".format(i)
study_id = storage.create_new_study(study_name=study_name)
if direction is None:
direction = generator.choice(
[StudyDirection.MINIMIZE, StudyDirection.MAXIMIZE])
storage.set_study_direction(study_id, direction)
best_trial = None
trials = {}
datetime_start = None
for j in range(n_trial):
trial = _generate_trial(generator)
trial.number = j
trial._trial_id = storage.create_new_trial(study_id, trial)
trials[trial._trial_id] = trial
if datetime_start is None:
datetime_start = trial.datetime_start
else:
datetime_start = min(datetime_start, trial.datetime_start)
if trial.state == TrialState.COMPLETE and trial.value is not None:
if best_trial is None:
best_trial = trial
else:
if direction == StudyDirection.MINIMIZE and trial.value < best_trial.value:
best_trial = trial
elif direction == StudyDirection.MAXIMIZE and best_trial.value < trial.value:
best_trial = trial
study_id_to_trials[study_id] = trials
study_id_to_summary[study_id] = StudySummary(
study_name=study_name,
direction=direction,
best_trial=best_trial,
user_attrs={},
system_attrs={},
n_trials=len(trials),
datetime_start=datetime_start,
study_id=study_id,
)
return study_id_to_summary, study_id_to_trials
def get_study_summary(storage: BaseStorage,
study_id: int) -> Optional[StudySummary]:
summaries = storage.get_all_study_summaries()
for summary in summaries:
if summary._study_id != study_id:
continue
return summary
return None
def is_heartbeat_enabled(storage: BaseStorage) -> bool:
"""Check whether the storage enables the heartbeat.
Returns:
:obj:`True` if the storage also inherits :class:`~optuna.storages._heartbeat.BaseHeartbeat`
and the return value of :meth:`~optuna.storages.BaseStorage.get_heartbeat_interval` is an
integer, otherwise :obj:`False`.
"""
return isinstance(
storage,
BaseHeartbeat) and storage.get_heartbeat_interval() is not None
def get_trials(storage: BaseStorage,
study_id: int,
ttl_seconds: int = 10) -> List[FrozenTrial]:
with trials_cache_lock:
trials = trials_cache.get(study_id, None)
last_fetched_at = trials_last_fetched_at.get(study_id, None)
if (trials is not None and last_fetched_at is not None
and datetime.now() - last_fetched_at < timedelta(ttl_seconds)):
return trials
trials = storage.get_all_trials(study_id)
with trials_cache_lock:
trials_last_fetched_at[study_id] = datetime.now()
trials_cache[study_id] = trials
return trials