Add support MLFlow plugin

[kumare3]

Signed-off-by: Ketan Umare [email protected]

TL;DR

MLFlow plugin for flytekit. Render mlflow metrics and parameters on the flyte decks.

Example:

from flytekit import task, workflow
from flytekitplugins.mlflow import mlflow_autolog
import mlflow
import tensorflow as tf


@task(disable_deck=False)
@mlflow_autolog(framework=mlflow.keras)
def train_model(epochs: int):
    # Refer to https://www.tensorflow.org/tutorials/keras/classification
    fashion_mnist = tf.keras.datasets.fashion_mnist
    (train_images, train_labels), (_, _) = fashion_mnist.load_data()
    train_images = train_images / 255.0

    model = tf.keras.Sequential([
        tf.keras.layers.Flatten(input_shape=(28, 28)),
        tf.keras.layers.Dense(128, activation='relu'),
        tf.keras.layers.Dense(10)
    ])

    model.compile(optimizer='adam',
                  loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
                  metrics=['accuracy'])
    model.fit(train_images, train_labels, epochs=epochs)


@workflow
def ml_pipeline(epochs: int):
    train_model(epochs=epochs)


if __name__ == "__main__":
    ml_pipeline()

Type

• Bug Fix
• Feature
• Plugin

Are all requirements met?

• Code completed
• Smoke tested
• Unit tests added
• Code documentation added
• Any pending items have an associated Issue

Complete description

Tracking Issue

https://github.com/flyteorg/flyte/issues/

[codecov]

Codecov Report

Merging #1274 (e3c49c0) into master (8c72082) will not change coverage.
The diff coverage is n/a.

@@           Coverage Diff           @@
##           master    #1274   +/-   ##
=======================================
  Coverage   69.18%   69.18%           
=======================================
  Files         299      299           
  Lines       27996    27996           
  Branches     2631     2631           
=======================================
  Hits        19368    19368           
  Misses       8114     8114           
  Partials      514      514           

Help us with your feedback. Take ten seconds to tell us how you rate us. Have a feature suggestion? Share it here.

[pingsutw]

cc @kumare3 @cosmicBboy Should we use flytesnack_development as experiment name, and execution_name as run name here? In mlflow, each experiment can have multiple runs (or executions), so the experiment here is like the project in Flyte.

[cosmicBboy]

re: experiment and run name conventions, we have to think about how this would work in practice:

• an experiment is a group of runs, where the run produces comparable artifacts
• a run is typically a training job of the same type of model with different configuration (hyperparameters and even data)

So a more realistic example would be:

@task(disable_deck=False)
@mlflow_autolog(framework=mlflow.keras)
def train_model(hyperparams: dict):
    ...

@dynamic
def tune_model(hyperparams_grid: List[dict]):
    for hyperparams in hyperparams_grid:
        train_model(hyperparams=hyperparams)

@workflow
def ml_pipeline(hyperparams_grid: List[int]):
    tune_model(hyperparams_grid=hyperparams_grid)

Ideally, these would be the requirements of an MLFlow plugin, which doesn't block this PR but just to note for a future iteration:

• Experiment name should be consistent across different runs of the same experiment, so probably either the workflow name, or an explicitly passed-in experiment name, defined at the top-level scope of the module mlflow_autolog(experiment_name=EXPERIMENT_NAME)
• The run name should be unique to the particular task execution, so that execution_name + task_node + <SOMETHING ELSE> would make sense, where <SOMETHING ELSE> would be anything else that would uniquely identify the task execution.

[pingsutw]

Update run name and set experiment name while running.

	local	remote
Run name	random string created by MLFlow	<execution_id>.<task_name>
Experiment name	local	<project_name>.<domain_name>.<workflow_name>

To write the MLFlow metric to Postgres, update the default-env-vars for the tasks

• propeller-config

    plugins:
      k8s:
        default-cpus: 100m
        default-env-vars:
        - FLYTE_AWS_ENDPOINT: http://minio.flyte.svc.cluster.local:9000
        - FLYTE_AWS_ACCESS_KEY_ID: minio
        - FLYTE_AWS_SECRET_ACCESS_KEY: miniostorage
        - MLFLOW_TRACKING_URI: postgresql+psycopg2://postgres:@postgres.flyte.svc.cluster.local:5432/flyteadmin

[cosmicBboy]

Let's add MLFlow to the plugins docs: https://raw.githubusercontent.com/flyteorg/flytekit/master/docs/source/plugins/index.rst

[cosmicBboy]

hmm, where are the flytekitplugins-* being installed in the docs build process?

I think we need to update https://github.com/flyteorg/flytekit/blob/master/doc-requirements.in with the mlflow plugin dependencies, the autolog function isn't showing up in the docs preview: https://flyte--1274.org.readthedocs.build/projects/flytekit/en/1274/plugins/mlflow.html

[cosmicBboy]

🚀