optuna_integration.MLflowCallback

class optuna_integration.MLflowCallback(tracking_uri=None, metric_name='value', create_experiment=True, mlflow_kwargs=None, tag_study_user_attrs=False, tag_trial_user_attrs=True)[source]

Callback to track Optuna trials with MLflow.

This callback adds relevant information that is tracked by Optuna to MLflow.

Example

Add MLflow callback to Optuna optimization.

Parameters:

  • tracking_uri (str | None) –

    The URI of the MLflow tracking server.

    Please refer to mlflow.set_tracking_uri for more details.

  • metric_name (str | Sequence[str]) – Name assigned to optimized metric. In case of multi-objective optimization, list of names can be passed. Those names will be assigned to metrics in the order returned by objective function. If single name is provided, or this argument is left to default value, it will be broadcasted to each objective with a number suffix in order returned by objective function e.g. two objectives and default metric name will be logged as value_0 and value_1. The number of metrics must be the same as the number of values an objective function returns.

  • create_experiment (bool) – When True, new MLflow experiment will be created for each optimization run, named after the Optuna study. Setting this argument to False lets user run optimization under existing experiment, set via mlflow.set_experiment, by passing experiment_id as one of mlflow_kwargs or under default MLflow experiment, when no additional arguments are passed. Note that this argument must be set to False when using Optuna with this callback within Databricks Notebook.

  • mlflow_kwargs (dict[str, Any] | None) –

    Set of arguments passed when initializing MLflow run. Please refer to MLflow API documentation for more details.

    Note

    nest_trials argument added in v2.3.0 is a part of mlflow_kwargs since v3.0.0. Anyone using nest_trials=True should migrate to mlflow_kwargs={"nested": True} to avoid raising TypeError.

  • tag_study_user_attrs (bool) – Flag indicating whether or not to add the study’s user attrs to the mlflow trial as tags. Please note that when this flag is set, key value pairs in user_attrs will supersede existing tags.

  • tag_trial_user_attrs (bool) – Flag indicating whether or not to add the trial’s user attrs to the mlflow trial as tags. Please note that when both trial and study user attributes are logged, the latter will supersede the former in case of a collision.

Warning

Deprecated in v4.9.0. This feature will be removed in the future. The removal of this feature is currently scheduled for v6.0.0, but this schedule is subject to change. See https://github.com/optuna/optuna/releases/tag/v4.9.0.

Methods

track_in_mlflow()

Decorator for using MLflow logging in the objective function.
track_in_mlflow()[source]

Decorator for using MLflow logging in the objective function.

This decorator enables the extension of MLflow logging provided by the callback.

All information logged in the decorated objective function will be added to the MLflow run for the trial created by the callback.

Example

Add additional logging to MLflow.

Returns:

Objective function with tracking to MLflow enabled.

Return type:

Callable

Note

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