`OptunaMLflow`

See also

Code documentation can be found here: code documentation page

OptunaMLflow is a wrapper to use Optuna and log to MLflow at the same time. If an exception occurs during communication with MLflow, it is caught and handled without interrupting the training process. The OptunaMLflow class is used as a decorator for Optuna objective functions. It looks like this:

from hpoflow import OptunaMLflow
import optuna

@OptunaMLflow()
def objective(trial):
    x = trial.suggest_float("x", -10, 10)
    return (x - 2) ** 2

study = optuna.create_study()
study.optimize(objective, n_trials=100)

If this decorator is applied the Optuna optuna.study.Study object is augmented. This augmentation entails logging information to Optuna and MLflow in parallel.

Autologging to MLflow

You can use Optuna and the Trial object as you are used to. The following is automatically logged in parallel:

the Optuna distributions and parameters (name and ranges)
the sampled hyperparameters
the objective result
exceptions with traceback
the direction of the study (MINIMIZE or MAXIMIZE)
hostname and process id

Manual Logging

In addition to what is logged automatically by Optuna usage, the following can be logged manually:

additonal metrics: log_metric() or log_metrics()
parameter: log_param() or log_params()
tags: set_tag() or set_tags()

Note

The manually logged information is stored by Optuna as a user attribute on the Trial object (also see set_user_attr()).

Logging Nested Runs

Sometimes you want to repeat a training several times with the same hyperparameters within a trial. This is the case, for example, when performing a cross-validation. It is possible to log the results of these repetitions as so-called nested runs on the MLflow side. To do this use the log_iter() method. It looks like this:

from hpoflow import OptunaMLflow
import numpy as np
import optuna

@OptunaMLflow()
def objective(trial):
    x = trial.suggest_uniform("x", -10, 10)

    results = []

    for i in range(7):  # simulate 7 fold cross-validation
        result = (x - 2) ** 2
        trial.log_iter({"fold_result": result}, i)  # call to log the fold as nested run
        results.append(result)

    result = np.mean(results)
    return result  # auto logging - no explicit call to log_metric needed

study = optuna.create_study()
study.optimize(objective, n_trials=100)

Note

Optuna does not support nested runs. That is why the results are aggregated into lists when they are stored as user attributes at Optuna.

Set MLflow Tracking Server URI

By passing tracking_uri to the constructor of OptunaMLflow you can set the MLflow tracking server URI (also see mlflow.set_tracking_uri()). The values can be:

not set or None: MLflow logs to the default locale folder ./mlruns or uses the MLFLOW_TRACKING_URI environment variable if it is available.
local file path, prefixed with file:/: Data is stored locally at the provided directory.
HTTP URI like https://my-tracking-server:5000
Databricks workspace, provided as the string databricks or, to use a Databricks CLI profile, databricks://<profileName>

Enforce no uncommitted GIT Changes

By passing enforce_clean_git=True to the constructor of OptunaMLflow you can check and enforce that the GIT repository has no uncommitted changes (see git.repo.base.Repo.is_dirty()). If there are uncommitted GIT changes an exception is raised. In this way, reproducibility of experiments is facilitated.

OptunaMLflow