helios.plugins.optuna.utils =========================== .. py:module:: helios.plugins.optuna.utils Classes ------- .. autoapisummary:: helios.plugins.optuna.utils.StudyArgs Functions --------- .. autoapisummary:: helios.plugins.optuna.utils.checkpoint_sampler helios.plugins.optuna.utils.restore_sampler helios.plugins.optuna.utils.create_or_load_study helios.plugins.optuna.utils.create_study_starting_from_trial Module Contents --------------- .. py:class:: StudyArgs Bases: :py:obj:`TypedDict` Type-hint class for the arguments of ``optuna.create_study``. See the `docs `__ for more information. :param storage: Database URL. If this argument is set to None, ``InMemoryStorage`` is used, and the ``Study`` will not be persistent. :param sampler: A sampler object that implements background algorithm for value suggestion. :param pruner: A pruner object that decides early stopping for unpromising trials. :param study_name: Study's name. :param direction: Direction of optimization. :param load_if_exists: flag to control the behaviour to handle conflict of study names. :param directions: A sequence of directions during multi-objective optimization. .. py:attribute:: storage :type: NotRequired[str | optuna.storages.BaseStorage | None] .. py:attribute:: sampler :type: NotRequired[optuna.samplers.BaseSampler | None] .. py:attribute:: pruner :type: NotRequired[optuna.pruners.BasePruner | None] .. py:attribute:: study_name :type: NotRequired[str | None] .. py:attribute:: direction :type: NotRequired[str | optuna.study.StudyDirection] .. py:attribute:: load_if_exists :type: NotRequired[bool] .. py:attribute:: directions :type: NotRequired[Sequence[str | optuna.study.StudyDirection] | None] .. py:function:: checkpoint_sampler(trial: optuna.Trial, chkpt_root: pathlib.Path) -> None 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. :param trial: the current trial. :param chkpt_root: the root where the checkpoints will be saved. .. py:function:: restore_sampler(chkpt_root: pathlib.Path, trial_id: int | None = None) -> optuna.samplers.BaseSampler | None Restore the sampler from a previously saved checkpoint. This function can be used in tandem with :py:func:`~helios.plugins.optuna.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. :param chkpt_root: the root where the checkpoints are stored. :param trial_id: the trial ID for which to restore the sampler. :returns: The restored sampler. .. py:function:: create_or_load_study(sampler_path: pathlib.Path | None = None, backup_study: bool = False, backup_samplers: bool = False, **study_args: Unpack[StudyArgs]) -> optuna.Study 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 :py:func:`~helios.plugins.optuna.restore_sampler` to ensure the study can be resumed properly. .. warning:: This function **requires** the following conditions to be true: #. The study uses a string to specify the ``RDBStorage`` as the argument to ``storage`` for ``optuna.create_study``. #. The storage string starts with ``sqlite:///``. #. ``load_if_exists`` is set to true. If you have been using :py:func:`~helios.plugins.optuna.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 ``_backup-#`` where ```` 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 :py:func:`~helios.plugins.optuna.create_study_starting_from_trial`. :param sampler_path: the path where the sampler checkpoints are stored. Defaults to None. :param backup_study: if True, the original study is backed up so it can be re-used later on. Defaults to false. :param backup_samplers: if True, a backup will be made of the samplers checkpoint folder (if given). Defaults to false. :param study_args: arguments for ``optuna.create_study``. .. py:function:: 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 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 :py:func:`~helios.plugins.optuna.resume_study` unable to properly resume. To accomplish this, the function will do the following: #. 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``. #. 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: #. It is called **before** the trials are started. #. The study uses ``RDBStorage`` in ``study_args``. #. ``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 ``_backup-#`` where ```` 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. :param trial_id: the ID of the last trial from which the study should resume. :param backup_study: if True, the original study is backed up so it can be re-used later on. :param sampler_path: the path where the sampler checkpoints are stored. :param study_args: arguments for ``optuna.create_study``. :returns: The study with trials up to ``last_trial_number``.