flytekitplugins.mlflow.mlflow_autolog#

flytekitplugins.mlflow.mlflow_autolog(fn=None, *, framework=<module 'mlflow.sklearn' from '/home/docs/checkouts/readthedocs.org/user_builds/flytekit/envs/latest/lib/python3.9/site-packages/mlflow/sklearn/__init__.py'>, 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 (Optional[str]) – The MLFlow experiment name. If not provided, uses the Flyte execution name.