ray.air.integrations.mlflow.MLflowLoggerCallback#

class ray.air.integrations.mlflow.MLflowLoggerCallback(tracking_uri: Optional[str] = None, *, registry_uri: Optional[str] = None, experiment_name: Optional[str] = None, tags: Optional[Dict] = None, tracking_token: Optional[str] = None, save_artifact: bool = False)[source]#

Bases: ray.tune.logger.logger.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.

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.

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)])

Methods

get_state()

Get the state of the callback.

log_trial_restore(trial)

Handle logging when a trial restores.

log_trial_save(trial)

Handle logging when a trial saves a checkpoint.

on_checkpoint(iteration, trials, trial, ...)

Called after a trial saved a checkpoint with Tune.

on_experiment_end(trials, **info)

Called after experiment is over and all trials have concluded.

on_step_begin(iteration, trials, **info)

Called at the start of each tuning loop step.

on_step_end(iteration, trials, **info)

Called at the end of each tuning loop step.

on_trial_recover(iteration, trials, trial, ...)

Called after a trial instance failed (errored) but the trial is scheduled for retry.

set_state(state)

Set the state of the callback.