ray.air.integrations.mlflow.MLflowLoggerCallback

class ray.air.integrations.mlflow.MLflowLoggerCallback(tracking_uri: str | None = None, *, registry_uri: str | None = None, experiment_name: str | None = None, tags: Dict | None = None, tracking_token: str | None = None, save_artifact: bool = False, log_params_on_trial_end: bool = False)

Bases: LoggerCallback

MLflow Logger to automatically log Tune results and config to MLflow.

MLflow (https://mlflow.org) Tracking is an open source library for recording and querying experiments. This Ray Tune LoggerCallback sends information (config parameters, training results & metrics, and artifacts) to MLflow for automatic experiment tracking.

Keep in mind that the callback will open an MLflow session on the driver and not on the trainable. Therefore, it is not possible to call MLflow functions like mlflow.log_figure() inside the trainable as there is no MLflow session on the trainable. For more fine grained control, use ray.air.integrations.mlflow.setup_mlflow().

Parameters:

• tracking_uri – The tracking URI for where to manage experiments and runs. This can either be a local file path or a remote server. This arg gets passed directly to mlflow initialization. When using Tune in a multi-node setting, make sure to set this to a remote server and not a local file path.

• registry_uri – The registry URI that gets passed directly to mlflow initialization.

• experiment_name – The experiment name to use for this Tune run. If the experiment with the name already exists with MLflow, it will be reused. If not, a new experiment will be created with that name.

• tags – An optional dictionary of string keys and values to set as tags on the run

• tracking_token – Tracking token used to authenticate with MLflow.

• save_artifact – If set to True, automatically save the entire contents of the Tune local_dir as an artifact to the corresponding run in MlFlow.

• log_params_on_trial_end – If set to True, log parameters to MLflow at the end of the trial instead of at the beginning

Example:

from ray.air.integrations.mlflow import MLflowLoggerCallback

tags = { "user_name" : "John",
         "git_commit_hash" : "abc123"}

tune.run(
    train_fn,
    config={
        # define search space here
        "parameter_1": tune.choice([1, 2, 3]),
        "parameter_2": tune.choice([4, 5, 6]),
    },
    callbacks=[MLflowLoggerCallback(
        experiment_name="experiment1",
        tags=tags,
        save_artifact=True,
        log_params_on_trial_end=True)])

Methods

get_state	Get the state of the callback.
log_trial_restore	Handle logging when a trial restores.
log_trial_save	Handle logging when a trial saves a checkpoint.
on_checkpoint	Called after a trial saved a checkpoint with Tune.
on_experiment_end	Called after experiment is over and all trials have concluded.
on_step_begin	Called at the start of each tuning loop step.
on_step_end	Called at the end of each tuning loop step.
on_trial_recover	Called after a trial instance failed (errored) but the trial is scheduled for retry.
set_state	Set the state of the callback.