flytekitplugins.mlflow.mlflow_autolog(fn=None, *, framework=<module 'mlflow.sklearn (Not loaded yet)'>, experiment_name=None)[source]

MLFlow decorator to enable autologging of training metrics.

This decorator can be used as a nested decorator for a @task and it will automatically enable mlflow autologging, for the given framework. By default autologging is enabled for sklearn.

def my_tensorflow_trainer():

One benefit of doing so is that the mlflow metrics are then rendered inline using FlyteDecks and can be viewed in jupyter notebook, as well as in hosted Flyte environment:

# jupyter notebook cell
with flytekit.new_context() as ctx:
    ctx.get_deck()  # IPython.display

When the task is called in a Flyte backend, the decorator starts a new MLFlow run using the Flyte execution name by default, or a user-provided experiment_name in the decorator.

  • fn – Function to generate autologs for.

  • framework – The mlflow module to use for autologging

  • experiment_name (str | None) – The MLFlow experiment name. If not provided, uses the Flyte execution name.