flytekitplugins.mlflow.mlflow_autolog¶

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.

@task
@mlflow_autolog(framework=mlflow.tensorflow)
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:
    my_tensorflow_trainer()
    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.

Parameters:

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.