helios.plugins.optuna.utils

Classes

StudyArgs

Type-hint class for the arguments of optuna.create_study.

Functions

checkpoint_sampler(→ None)

Create a checkpoint with the state of the sampler.

restore_sampler(→ optuna.samplers.BaseSampler | None)

Restore the sampler from a previously saved checkpoint.

create_or_load_study(→ optuna.Study)

Start or continue a study.

create_study_starting_from_trial(→ optuna.Study)

Resume a study from a specific trial number.

Module Contents

class helios.plugins.optuna.utils.StudyArgs[source]

Bases: TypedDict

Type-hint class for the arguments of optuna.create_study.

See the docs for more information.

Parameters:

  • storage – Database URL. If this argument is set to None, InMemoryStorage is used, and the Study will not be persistent.

  • sampler – A sampler object that implements background algorithm for value suggestion.

  • pruner – A pruner object that decides early stopping for unpromising trials.

  • study_name – Study’s name.

  • direction – Direction of optimization.

  • load_if_exists – flag to control the behaviour to handle conflict of study names.

  • directions – A sequence of directions during multi-objective optimization.

storage: NotRequired[str | optuna.storages.BaseStorage | None]
sampler: NotRequired[optuna.samplers.BaseSampler | None]
pruner: NotRequired[optuna.pruners.BasePruner | None]
study_name: NotRequired[str | None]
direction: NotRequired[str | optuna.study.StudyDirection]
load_if_exists: NotRequired[bool]
directions: NotRequired[Sequence[str | optuna.study.StudyDirection] | None]
helios.plugins.optuna.utils.checkpoint_sampler(trial: optuna.Trial, chkpt_root: pathlib.Path) → None[source]

Create a checkpoint with the state of the sampler.

This function can be used to ensure that if a study is restarted, the state of the sampler is recovered so trials can be reproducible. The function will automatically create a checkpoint using torch.save.

Note

It is recommended that this function be called at the start and end of the objective function to ensure the state of the sampler is correctly stored. This ensures that if a trial fails before it returns, then the sampler can be restored and the same values can be retrieved. Alternatively, if the trial completes and you wish to resume the study, then the sampler can be restored to the state it would’ve had if the study hadn’t been stopped.

Parameters:

  • trial – the current trial.

  • chkpt_root – the root where the checkpoints will be saved.

helios.plugins.optuna.utils.restore_sampler(chkpt_root: pathlib.Path, trial_id: int | None = None) → optuna.samplers.BaseSampler | None[source]

Restore the sampler from a previously saved checkpoint.

This function can be used in tandem with checkpoint_sampler() to ensure that the last checkpoint is loaded and the correct state is restored for the sampler. This function needs to be called before optuna.create_study is called.

If you wish to restore the sampler for a specific trial, you can use trial_id to restore it.

Parameters:

  • chkpt_root – the root where the checkpoints are stored.

  • trial_id – the trial ID for which to restore the sampler.

Returns:

The restored sampler.

helios.plugins.optuna.utils.create_or_load_study(sampler_path: pathlib.Path | None = None, backup_study: bool = False, backup_samplers: bool = False, **study_args: Unpack[StudyArgs]) → optuna.Study[source]

Start or continue a study.

This function can be used to either start a new study or resume one that has stopped. This leverages the existing built-in checkpoint system from Helios as well as restore_sampler() to ensure the study can be resumed properly.

Warning

This function requires the following conditions to be true:
  1. The study uses a string to specify the RDBStorage as the argument to

    storage for optuna.create_study.

  2. The storage string starts with sqlite:///.

  3. load_if_exists is set to true.

If you have been using checkpoint_sampler() and wish to restore the sampler upon restarting the study, you can provide the root where the checkpoints are stored by passing in sampler_path.

You can use backup_study to create a backup of the original study. If there is one, the backup will be called <study-name>_backup-# where <study-name> is the name of the database of the original study and # is an incremental number starting at 0.

Note

This function does not retry failed trials, nor does it continue trials that were stopped. It will simply continue with the next (new) trial as determined by Optuna. To restart a study from a specific trial, please use create_study_starting_from_trial().

Parameters:

  • sampler_path – the path where the sampler checkpoints are stored. Defaults to None.

  • backup_study – if True, the original study is backed up so it can be re-used later on. Defaults to false.

  • backup_samplers – if True, a backup will be made of the samplers checkpoint folder (if given). Defaults to false.

  • study_args – arguments for optuna.create_study.

helios.plugins.optuna.utils.create_study_starting_from_trial(trial_id: int, backup_study: bool = True, backup_samplers: bool = True, sampler_path: pathlib.Path | None = None, **study_args: Unpack[StudyArgs]) → optuna.Study[source]

Resume a study from a specific trial number.

The goal of this function is to allow studies to resume from a specific trial number, effectively discarding all trials after that point. This can be used to recover from certain types of unexpected errors that result in trials being invalid but still marked as completed, thereby making resume_study() unable to properly resume. To accomplish this, the function will do the following:

  1. Grab all the trials from the study created by the given study_args and grab all

    trials whose numbers are less than or equal to last_trial_number.

  2. Create a new study with the same name and storage. This new study will get all of

    the trials that were obtained in the previous step.

If sampler_path is provided, then the samplers for trials that occur after trial_id will be pruned. The sampler that corresponds to trial_id will be preserved so it can be restored in to the new study. If you don’t wish this to happen, simply pass in the sampler in study_args and set sampler_path to None.

Note

Trials are 0-indexed. This means that trial N maps to ID N - 1.

Warning

This function requires the following conditions to be true:

  1. It is called before the trials are started.

  2. The study uses RDBStorage in study_args.

  3. load_if_exists is set to True in study_args.

By default, the original study and samplers checkpoint folder will be backed up (assuming they exist). The backups are named <name>_backup-# where <name> is either the name of the database of the original study or the sampler folder, and # is an incremental number starting at 0. This behaviour can be disabled by setting backup_study or backup_samplers to false.

In addition, this function can also remove the saved samplers from the discarded trials if sampler_root is provided.

Parameters:

  • trial_id – the ID of the last trial from which the study should resume.

  • backup_study – if True, the original study is backed up so it can be re-used later on.

  • sampler_path – the path where the sampler checkpoints are stored.

  • study_args – arguments for optuna.create_study.

Returns:

The study with trials up to last_trial_number.