optuna.integration.DaskStorage

class optuna.integration.DaskStorage(storage=None, name=None, client=None, register=True)[source]

Dask-compatible storage class.

This storage class wraps a Optuna storage class (e.g. Optuna’s in-memory or sqlite storage) and is used to run optimization trials in parallel on a Dask cluster. The underlying Optuna storage object lives on the cluster’s scheduler and any method calls on the DaskStorage instance results in the same method being called on the underlying Optuna storage object.

See this example for how to use DaskStorage to extend Optuna’s in-memory storage class to run across multiple processes.

Parameters:

storage (None | str | BaseStorage) – Optuna storage url to use for underlying Optuna storage class to wrap (e.g. None for in-memory storage, sqlite:///example.db for SQLite storage). Defaults to None.
name (str | None) – Unique identifier for the Dask storage class. Specifying a custom name can sometimes be useful for logging or debugging. If None is provided, a random name will be automatically generated.
client (distributed.Client | None) – Dask Client to connect to. If not provided, will attempt to find an existing Client.
register (bool) – Whether or not to register this storage instance with the cluster scheduler. Most common usage of this storage class will not need to specify this argument. Defaults to True.

Note

Added in v3.1.0 as an experimental feature. The interface may change in newer versions without prior notice. See https://github.com/optuna/optuna/releases/tag/v3.1.0.

Methods

`check_trial_is_updatable`(trial_id, trial_state)	Check whether a trial state is updatable.
`create_new_study`(directions[, study_name])	Create a new study from a name.
`create_new_trial`(study_id[, template_trial])	Create and add a new trial to a study.
`delete_study`(study_id)	Delete a study.
`get_all_studies`()	Read a list of `FrozenStudy` objects.
`get_all_trials`(study_id[, deepcopy, states])	Read all trials in a study.
`get_base_storage`()	Retrieve underlying Optuna storage instance from the scheduler.
`get_best_trial`(study_id)	Return the trial with the best value in a study.
`get_n_trials`(study_id[, state])	Count the number of trials in a study.
`get_study_directions`(study_id)	Read whether a study maximizes or minimizes an objective.
`get_study_id_from_name`(study_name)	Read the ID of a study.
`get_study_name_from_id`(study_id)	Read the study name of a study.
`get_study_system_attrs`(study_id)	Read the optuna-internal attributes of a study.
`get_study_user_attrs`(study_id)	Read the user-defined attributes of a study.
`get_trial`(trial_id)	Read a trial.
`get_trial_id_from_study_id_trial_number`(...)	Read the trial ID of a trial.
`get_trial_number_from_id`(trial_id)	Read the trial number of a trial.
`get_trial_param`(trial_id, param_name)	Read the parameter of a trial.
`get_trial_params`(trial_id)	Read the parameter dictionary of a trial.
`get_trial_system_attrs`(trial_id)	Read the optuna-internal attributes of a trial.
`get_trial_user_attrs`(trial_id)	Read the user-defined attributes of a trial.
`remove_session`()	Clean up all connections to a database.
`set_study_system_attr`(study_id, key, value)	Register an optuna-internal attribute to a study.
`set_study_user_attr`(study_id, key, value)	Register a user-defined attribute to a study.
`set_trial_intermediate_value`(trial_id, step, ...)	Report an intermediate value of an objective function.
`set_trial_param`(trial_id, param_name, ...)	Set a parameter to a trial.
`set_trial_state_values`(trial_id, state[, values])	Update the state and values of a trial.
`set_trial_system_attr`(trial_id, key, value)	Set an optuna-internal attribute to a trial.
`set_trial_user_attr`(trial_id, key, value)	Set a user-defined attribute to a trial.

Attributes

client

check_trial_is_updatable(trial_id, trial_state)

Check whether a trial state is updatable.

Parameters:

trial_id (int) – ID of the trial. Only used for an error message.
trial_state (TrialState) – Trial state to check.

Raises:

RuntimeError – If the trial is already finished.

Return type:

None

create_new_study(directions, study_name=None)[source]

Create a new study from a name.

If no name is specified, the storage class generates a name. The returned study ID is unique among all current and deleted studies.

Parameters:

directions (Sequence[StudyDirection]) – A sequence of direction whose element is either MAXIMIZE or MINIMIZE.
study_name (str | None) – Name of the new study to create.

Returns:

ID of the created study.

Raises:

optuna.exceptions.DuplicatedStudyError – If a study with the same study_name already exists.

Return type:

int

create_new_trial(study_id, template_trial=None)[source]

Create and add a new trial to a study.

The returned trial ID is unique among all current and deleted trials.

Parameters:

study_id (int) – ID of the study.
template_trial (FrozenTrial | None) – Template FrozenTrial with default user-attributes, system-attributes, intermediate-values, and a state.

Returns:

ID of the created trial.

Raises:

KeyError – If no study with the matching study_id exists.

Return type:

int

delete_study(study_id)[source]

Delete a study.

Parameters:: study_id (int) – ID of the study.
Raises:: KeyError – If no study with the matching study_id exists.
Return type:: None

get_all_studies()[source]

Read a list of FrozenStudy objects.

Returns:: A list of FrozenStudy objects, sorted by study_id.
Return type:: List[FrozenStudy]

get_all_trials(study_id, deepcopy=True, states=None)[source]

Read all trials in a study.

Parameters:

study_id (int) – ID of the study.
deepcopy (bool) – Whether to copy the list of trials before returning. Set to True if you intend to update the list or elements of the list.
states (Container[TrialState] | None) – Trial states to filter on. If None, include all states.

Returns:

List of trials in the study, sorted by trial_id.

Raises:

KeyError – If no study with the matching study_id exists.

Return type:

List[FrozenTrial]

get_base_storage()[source]

Retrieve underlying Optuna storage instance from the scheduler.

This is a convenience method to extract the Optuna storage instance stored on the Dask scheduler process to the local Python process.

Return type:: BaseStorage

get_best_trial(study_id)

Return the trial with the best value in a study.

This method is valid only during single-objective optimization.

Parameters:

study_id (int) – ID of the study.

Returns:

The trial with the best objective value among all finished trials in the study.

Raises:

KeyError – If no study with the matching study_id exists.
RuntimeError – If the study has more than one direction.
ValueError – If no trials have been completed.

Return type:

FrozenTrial

get_n_trials(study_id, state=None)[source]

Count the number of trials in a study.

Parameters:

study_id (int) – ID of the study.
state (Tuple[TrialState, ...] | TrialState | None) – Trial states to filter on. If None, include all states.

Returns:

Number of trials in the study.

Raises:

KeyError – If no study with the matching study_id exists.

Return type:

int

get_study_directions(study_id)[source]

Read whether a study maximizes or minimizes an objective.

Parameters:: study_id (int) – ID of a study.
Returns:: Optimization directions list of the study.
Raises:: KeyError – If no study with the matching study_id exists.
Return type:: List[StudyDirection]

get_study_id_from_name(study_name)[source]

Read the ID of a study.

Parameters:: study_name (str) – Name of the study.
Returns:: ID of the study.
Raises:: KeyError – If no study with the matching study_name exists.
Return type:: int

get_study_name_from_id(study_id)[source]

Read the study name of a study.

Parameters:: study_id (int) – ID of the study.
Returns:: Name of the study.
Raises:: KeyError – If no study with the matching study_id exists.
Return type:: str

get_study_system_attrs(study_id)[source]

Read the optuna-internal attributes of a study.

Parameters:: study_id (int) – ID of the study.
Returns:: Dictionary with the optuna-internal attributes of the study.
Raises:: KeyError – If no study with the matching study_id exists.
Return type:: Dict[str, Any]

get_study_user_attrs(study_id)[source]

Read the user-defined attributes of a study.

Parameters:: study_id (int) – ID of the study.
Returns:: Dictionary with the user attributes of the study.
Raises:: KeyError – If no study with the matching study_id exists.
Return type:: Dict[str, Any]

get_trial(trial_id)[source]

Read a trial.

Parameters:: trial_id (int) – ID of the trial.
Returns:: Trial with a matching trial ID.
Raises:: KeyError – If no trial with the matching trial_id exists.
Return type:: FrozenTrial

get_trial_id_from_study_id_trial_number(study_id, trial_number)[source]

Read the trial ID of a trial.

Parameters:

study_id (int) – ID of the study.
trial_number (int) – Number of the trial.

Returns:

ID of the trial.

Raises:

KeyError – If no trial with the matching study_id and trial_number exists.

Return type:

int

get_trial_number_from_id(trial_id)[source]

Read the trial number of a trial.

Note

The trial number is only unique within a study, and is sequential.

Parameters:: trial_id (int) – ID of the trial.
Returns:: Number of the trial.
Raises:: KeyError – If no trial with the matching trial_id exists.
Return type:: int

get_trial_param(trial_id, param_name)[source]

Read the parameter of a trial.

Parameters:

trial_id (int) – ID of the trial.
param_name (str) – Name of the parameter.

Returns:

Internal representation of the parameter.

Raises:

KeyError – If no trial with the matching trial_id exists. If no such parameter exists.

Return type:

float

get_trial_params(trial_id)

Read the parameter dictionary of a trial.

Parameters:: trial_id (int) – ID of the trial.
Returns:: Dictionary of a parameters. Keys are parameter names and values are internal representations of the parameter values.
Raises:: KeyError – If no trial with the matching trial_id exists.
Return type:: Dict[str, Any]

get_trial_system_attrs(trial_id)

Read the optuna-internal attributes of a trial.

Parameters:: trial_id (int) – ID of the trial.
Returns:: Dictionary with the optuna-internal attributes of the trial.
Raises:: KeyError – If no trial with the matching trial_id exists.
Return type:: Dict[str, Any]

get_trial_user_attrs(trial_id)

Read the user-defined attributes of a trial.

Parameters:: trial_id (int) – ID of the trial.
Returns:: Dictionary with the user-defined attributes of the trial.
Raises:: KeyError – If no trial with the matching trial_id exists.
Return type:: Dict[str, Any]

remove_session()

Clean up all connections to a database.

Return type:: None

set_study_system_attr(study_id, key, value)[source]

Register an optuna-internal attribute to a study.

This method overwrites any existing attribute.

Parameters:

study_id (int) – ID of the study.
key (str) – Attribute key.
value (Any) – Attribute value. It should be JSON serializable.

Raises:

KeyError – If no study with the matching study_id exists.

Return type:

None

set_study_user_attr(study_id, key, value)[source]

Register a user-defined attribute to a study.

This method overwrites any existing attribute.

Parameters:

study_id (int) – ID of the study.
key (str) – Attribute key.
value (Any) – Attribute value. It should be JSON serializable.

Raises:

KeyError – If no study with the matching study_id exists.

Return type:

None

set_trial_intermediate_value(trial_id, step, intermediate_value)[source]

Report an intermediate value of an objective function.

This method overwrites any existing intermediate value associated with the given step.

Parameters:

trial_id (int) – ID of the trial.
step (int) – Step of the trial (e.g., the epoch when training a neural network).
intermediate_value (float) – Intermediate value corresponding to the step.

Raises:

KeyError – If no trial with the matching trial_id exists.
RuntimeError – If the trial is already finished.

Return type:

None

set_trial_param(trial_id, param_name, param_value_internal, distribution)[source]

Set a parameter to a trial.

Parameters:

trial_id (int) – ID of the trial.
param_name (str) – Name of the parameter.
param_value_internal (float) – Internal representation of the parameter value.
distribution (BaseDistribution) – Sampled distribution of the parameter.

Raises:

KeyError – If no trial with the matching trial_id exists.
RuntimeError – If the trial is already finished.

Return type:

None

set_trial_state_values(trial_id, state, values=None)[source]

Update the state and values of a trial.

Set return values of an objective function to values argument. If values argument is not None, this method overwrites any existing trial values.

Parameters:

trial_id (int) – ID of the trial.
state (TrialState) – New state of the trial.
values (Sequence[float] | None) – Values of the objective function.

Returns:

True if the state is successfully updated. False if the state is kept the same. The latter happens when this method tries to update the state of RUNNING trial to RUNNING.

Raises:

KeyError – If no trial with the matching trial_id exists.
RuntimeError – If the trial is already finished.

Return type:

bool

set_trial_system_attr(trial_id, key, value)[source]

Set an optuna-internal attribute to a trial.

This method overwrites any existing attribute.

Parameters:

trial_id (int) – ID of the trial.
key (str) – Attribute key.
value (Mapping[str, Mapping[str, JSONSerializable] | Sequence[JSONSerializable] | str | int | float | bool | None] | Sequence[Mapping[str, JSONSerializable] | Sequence[JSONSerializable] | str | int | float | bool | None] | str | int | float | bool | None) – Attribute value. It should be JSON serializable.

Raises:

KeyError – If no trial with the matching trial_id exists.
RuntimeError – If the trial is already finished.

Return type:

None

set_trial_user_attr(trial_id, key, value)[source]

Set a user-defined attribute to a trial.

This method overwrites any existing attribute.

Parameters:

trial_id (int) – ID of the trial.
key (str) – Attribute key.
value (Any) – Attribute value. It should be JSON serializable.

Raises:

KeyError – If no trial with the matching trial_id exists.
RuntimeError – If the trial is already finished.

Return type:

None