chore: import upstream snapshot with attribution

examples/multistep_workflow/MLproject

@@ -0,0 +1,37 @@
+name: multistep_example
+
+python_env: python_env.yaml
+
+entry_points:
+  load_raw_data:
+    command: "python load_raw_data.py"
+
+  etl_data:
+    parameters:
+      ratings_csv: path
+      max_row_limit: {type: int, default: 100000}
+    command: "python etl_data.py --ratings-csv {ratings_csv} --max-row-limit {max_row_limit}"
+
+  als:
+    parameters:
+      ratings_data: path
+      max_iter: {type: int, default: 10}
+      reg_param: {type: float, default: 0.1}
+      rank: {type: int, default: 12}
+    command: "python als.py --ratings-data {ratings_data} --max-iter {max_iter} --reg-param {reg_param} --rank {rank}"
+
+  train_keras:
+    parameters:
+      ratings_data: path
+      als_model_uri: string
+      hidden_units: {type: int, default: 20}
+    command: "python train_keras.py --ratings-data {ratings_data} --als-model-uri {als_model_uri} --hidden-units {hidden_units}"
+
+  main:
+    parameters:
+      als_max_iter: {type: int, default: 10}
+      keras_hidden_units: {type: int, default: 20}
+      max_row_limit: {type: int, default: 100000}
+    command: "python main.py --als-max-iter {als_max_iter} --keras-hidden-units {keras_hidden_units}
+                             --max-row-limit {max_row_limit}"
+

examples/multistep_workflow/README.rst

@@ -0,0 +1,64 @@
+Multistep Workflow Example
+--------------------------
+This MLproject aims to be a fully self-contained example of how to
+chain together multiple different MLflow runs which each encapsulate
+a transformation or training step, allowing a clear definition of the
+interface between the steps, as well as allowing for caching and reuse
+of the intermediate results.
+
+At a high level, our goal is to predict users' ratings of movie given
+a history of their ratings for other movies. This example is based
+on `this webinar <https://databricks.com/blog/2018/07/13/scalable-end-to-end-deep-learning-using-tensorflow-and-databricks-on-demand-webinar-and-faq-now-available.html>`_
+by @brookewenig and @smurching.
+
+.. image:: ../../docs/source/_static/images/tutorial-multistep-workflow.png?raw=true
+
+There are four steps to this workflow:
+
+- **load_raw_data.py**: Downloads the MovieLens dataset
+  (a set of triples of user id, movie id, and rating) as a CSV and puts
+  it into the artifact store.
+
+- **etl_data.py**: Converts the MovieLens CSV from the
+  previous step into Parquet, dropping unnecessary columns along the way.
+  This reduces the input size from 500 MB to 49 MB, and allows columnar
+  access of the data.
+
+- **als.py**: Runs Alternating Least Squares for collaborative
+  filtering on the Parquet version of MovieLens to estimate the
+  movieFactors and userFactors. This produces a relatively accurate estimator.
+
+- **train_keras.py**: Trains a neural network on the
+  original data, supplemented by the ALS movie/userFactors -- we hope
+  this can improve upon the ALS estimations.
+
+While we can run each of these steps manually, here we have a driver
+run, defined as **main** (main.py). This run will run
+the steps in order, passing the results of one to the next.
+Additionally, this run will attempt to determine if a sub-run has
+already been executed successfully with the same parameters and, if so,
+reuse the cached results.
+
+Running this Example
+^^^^^^^^^^^^^^^^^^^^
+In order for the multistep workflow to find the other steps, you must
+execute ``mlflow run`` from this directory. So, in order to find out if
+the Keras model does in fact improve upon the ALS model, you can simply
+run:
+
+.. code-block:: bash
+
+    cd examples/multistep_workflow
+    mlflow run .
+
+
+This downloads and transforms the MovieLens dataset, trains an ALS
+model, and then trains a Keras model -- you can compare the results by
+using ``mlflow server``.
+
+You can also try changing the number of ALS iterations or Keras hidden
+units:
+
+.. code-block:: bash
+
+    mlflow run . -P als_max_iter=20 -P keras_hidden_units=50

examples/multistep_workflow/als.py

@@ -0,0 +1,70 @@
+"""
+Trains an Alternating Least Squares (ALS) model for user/movie ratings.
+The input is a Parquet ratings dataset (see etl_data.py), and we output
+an mlflow artifact called 'als-model'.
+"""
+
+import click
+import pyspark
+from pyspark.ml import Pipeline
+from pyspark.ml.evaluation import RegressionEvaluator
+from pyspark.ml.recommendation import ALS
+
+import mlflow
+import mlflow.spark
+
+
+@click.command()
+@click.option("--ratings-data")
+@click.option("--split-prop", default=0.8, type=float)
+@click.option("--max-iter", default=10, type=int)
+@click.option("--reg-param", default=0.1, type=float)
+@click.option("--rank", default=12, type=int)
+@click.option("--cold-start-strategy", default="drop")
+def train_als(ratings_data, split_prop, max_iter, reg_param, rank, cold_start_strategy):
+    seed = 42
+
+    with pyspark.sql.SparkSession.builder.getOrCreate() as spark:
+        ratings_df = spark.read.parquet(ratings_data)
+        (training_df, test_df) = ratings_df.randomSplit([split_prop, 1 - split_prop], seed=seed)
+        training_df.cache()
+        test_df.cache()
+
+        mlflow.log_metric("training_nrows", training_df.count())
+        mlflow.log_metric("test_nrows", test_df.count())
+
+        print(f"Training: {training_df.count()}, test: {test_df.count()}")
+
+        als = (
+            ALS()
+            .setUserCol("userId")
+            .setItemCol("movieId")
+            .setRatingCol("rating")
+            .setPredictionCol("predictions")
+            .setMaxIter(max_iter)
+            .setSeed(seed)
+            .setRegParam(reg_param)
+            .setColdStartStrategy(cold_start_strategy)
+            .setRank(rank)
+        )
+
+        als_model = Pipeline(stages=[als]).fit(training_df)
+
+        reg_eval = RegressionEvaluator(
+            predictionCol="predictions", labelCol="rating", metricName="mse"
+        )
+
+        predicted_test_dF = als_model.transform(test_df)
+
+        test_mse = reg_eval.evaluate(predicted_test_dF)
+        train_mse = reg_eval.evaluate(als_model.transform(training_df))
+
+        print(f"The model had a MSE on the test set of {test_mse}")
+        print(f"The model had a MSE on the (train) set of {train_mse}")
+        mlflow.log_metric("test_mse", test_mse)
+        mlflow.log_metric("train_mse", train_mse)
+        mlflow.spark.log_model(als_model, artifact_path="als-model")
+
+
+if __name__ == "__main__":
+    train_als()

examples/multistep_workflow/etl_data.py

@@ -0,0 +1,44 @@
+"""
+Converts the raw CSV form to a Parquet form with just the columns we want
+"""
+
+import os
+import tempfile
+
+import click
+import pyspark
+
+import mlflow
+
+
+@click.command(
+    help="Given a CSV file (see load_raw_data), transforms it into Parquet "
+    "in an mlflow artifact called 'ratings-parquet-dir'"
+)
+@click.option("--ratings-csv")
+@click.option(
+    "--max-row-limit", default=10000, help="Limit the data size to run comfortably on a laptop."
+)
+def etl_data(ratings_csv, max_row_limit):
+    with mlflow.start_run():
+        tmpdir = tempfile.mkdtemp()
+        ratings_parquet_dir = os.path.join(tmpdir, "ratings-parquet")
+        print(f"Converting ratings CSV {ratings_csv} to Parquet {ratings_parquet_dir}")
+        with pyspark.sql.SparkSession.builder.getOrCreate() as spark:
+            ratings_df = (
+                spark.read
+                .option("header", "true")
+                .option("inferSchema", "true")
+                .csv(ratings_csv)
+                .drop("timestamp")
+            )  # Drop unused column
+            ratings_df.show()
+            if max_row_limit != -1:
+                ratings_df = ratings_df.limit(max_row_limit)
+            ratings_df.write.parquet(ratings_parquet_dir)
+            print(f"Uploading Parquet ratings: {ratings_parquet_dir}")
+            mlflow.log_artifacts(ratings_parquet_dir, "ratings-parquet-dir")
+
+
+if __name__ == "__main__":
+    etl_data()

examples/multistep_workflow/load_raw_data.py

@@ -0,0 +1,43 @@
+"""
+Downloads the MovieLens dataset and saves it as an artifact
+"""
+
+import os
+import tempfile
+import zipfile
+
+import click
+import requests
+
+import mlflow
+
+
+@click.command(
+    help="Downloads the MovieLens dataset and saves it as an mlflow artifact "
+    "called 'ratings-csv-dir'."
+)
+@click.option("--url", default="http://files.grouplens.org/datasets/movielens/ml-20m.zip")
+def load_raw_data(url):
+    with mlflow.start_run():
+        local_dir = tempfile.mkdtemp()
+        local_filename = os.path.join(local_dir, "ml-20m.zip")
+        print(f"Downloading {url} to {local_filename}")
+        r = requests.get(url, stream=True)
+        with open(local_filename, "wb") as f:
+            for chunk in r.iter_content(chunk_size=1024):
+                if chunk:  # filter out keep-alive new chunks
+                    f.write(chunk)
+
+        extracted_dir = os.path.join(local_dir, "ml-20m")
+        print(f"Extracting {local_filename} into {extracted_dir}")
+        with zipfile.ZipFile(local_filename, "r") as zip_ref:
+            zip_ref.extractall(local_dir)
+
+        ratings_file = os.path.join(extracted_dir, "ratings.csv")
+
+        print(f"Uploading ratings: {ratings_file}")
+        mlflow.log_artifact(ratings_file, "ratings-csv-dir")
+
+
+if __name__ == "__main__":
+    load_raw_data()

examples/multistep_workflow/main.py

@@ -0,0 +1,107 @@
+"""
+Downloads the MovieLens dataset, ETLs it into Parquet, trains an
+ALS model, and uses the ALS model to train a Keras neural network.
+
+See README.md for more details.
+"""
+
+import os
+
+import click
+
+import mlflow
+from mlflow.entities import RunStatus
+from mlflow.tracking import MlflowClient
+from mlflow.tracking.fluent import _get_experiment_id
+from mlflow.utils import mlflow_tags
+from mlflow.utils.logging_utils import eprint
+
+
+def _already_ran(entry_point_name, parameters, git_commit, experiment_id=None):
+    """Best-effort detection of if a run with the given entrypoint name,
+    parameters, and experiment id already ran. The run must have completed
+    successfully and have at least the parameters provided.
+    """
+    experiment_id = experiment_id if experiment_id is not None else _get_experiment_id()
+    client = MlflowClient()
+    all_runs = reversed(client.search_runs([experiment_id]))
+    for run in all_runs:
+        tags = run.data.tags
+        if tags.get(mlflow_tags.MLFLOW_PROJECT_ENTRY_POINT, None) != entry_point_name:
+            continue
+        match_failed = False
+        for param_key, param_value in parameters.items():
+            run_value = run.data.params.get(param_key)
+            if run_value != param_value:
+                match_failed = True
+                break
+        if match_failed:
+            continue
+
+        if run.info.to_proto().status != RunStatus.FINISHED:
+            eprint(
+                ("Run matched, but is not FINISHED, so skipping (run_id={}, status={})").format(
+                    run.info.run_id, run.info.status
+                )
+            )
+            continue
+
+        previous_version = tags.get(mlflow_tags.MLFLOW_GIT_COMMIT, None)
+        if git_commit != previous_version:
+            eprint(
+                "Run matched, but has a different source version, so skipping "
+                f"(found={previous_version}, expected={git_commit})"
+            )
+            continue
+        return client.get_run(run.info.run_id)
+    eprint("No matching run has been found.")
+    return None
+
+
+# TODO(aaron): This is not great because it doesn't account for:
+# - changes in code
+# - changes in dependent steps
+def _get_or_run(entrypoint, parameters, git_commit, use_cache=True):
+    existing_run = _already_ran(entrypoint, parameters, git_commit)
+    if use_cache and existing_run:
+        print(f"Found existing run for entrypoint={entrypoint} and parameters={parameters}")
+        return existing_run
+    print(f"Launching new run for entrypoint={entrypoint} and parameters={parameters}")
+    submitted_run = mlflow.run(".", entrypoint, parameters=parameters, env_manager="local")
+    return MlflowClient().get_run(submitted_run.run_id)
+
+
+@click.command()
+@click.option("--als-max-iter", default=10, type=int)
+@click.option("--keras-hidden-units", default=20, type=int)
+@click.option("--max-row-limit", default=100000, type=int)
+def workflow(als_max_iter, keras_hidden_units, max_row_limit):
+    # Note: The entrypoint names are defined in MLproject. The artifact directories
+    # are documented by each step's .py file.
+    with mlflow.start_run() as active_run:
+        os.environ["SPARK_CONF_DIR"] = os.path.abspath(".")
+        git_commit = active_run.data.tags.get(mlflow_tags.MLFLOW_GIT_COMMIT)
+        load_raw_data_run = _get_or_run("load_raw_data", {}, git_commit)
+        ratings_csv_uri = os.path.join(load_raw_data_run.info.artifact_uri, "ratings-csv-dir")
+        etl_data_run = _get_or_run(
+            "etl_data", {"ratings_csv": ratings_csv_uri, "max_row_limit": max_row_limit}, git_commit
+        )
+        ratings_parquet_uri = os.path.join(etl_data_run.info.artifact_uri, "ratings-parquet-dir")
+
+        # We specify a spark-defaults.conf to override the default driver memory. ALS requires
+        # significant memory. The driver memory property cannot be set by the application itself.
+        als_run = _get_or_run(
+            "als", {"ratings_data": ratings_parquet_uri, "max_iter": str(als_max_iter)}, git_commit
+        )
+        als_model_uri = os.path.join(als_run.info.artifact_uri, "als-model")
+
+        keras_params = {
+            "ratings_data": ratings_parquet_uri,
+            "als_model_uri": als_model_uri,
+            "hidden_units": keras_hidden_units,
+        }
+        _get_or_run("train_keras", keras_params, git_commit, use_cache=False)
+
+
+if __name__ == "__main__":
+    workflow()

examples/multistep_workflow/python_env.yaml

@@ -0,0 +1,9 @@
+build_dependencies:
+  - pip
+dependencies:
+  - tensorflow==1.15.2
+  - keras==2.2.4
+  - mlflow>=1.0
+  - pyspark
+  - requests
+  - click

examples/multistep_workflow/spark-defaults.conf

@@ -0,0 +1 @@
+spark.driver.memory 8g

examples/multistep_workflow/train_keras.py

@@ -0,0 +1,117 @@
+"""
+Trains a Keras model for user/movie ratings. The input is a Parquet
+ratings dataset (see etl_data.py) and an ALS model (see als.py), which we
+will use to supplement our input and train using.
+"""
+
+from itertools import chain
+
+import click
+import numpy as np
+import pandas as pd
+import pyspark
+import tensorflow as tf
+from pyspark.sql.functions import col, udf
+from pyspark.sql.types import ArrayType, FloatType
+from tensorflow import keras
+from tensorflow.keras.callbacks import EarlyStopping
+from tensorflow.keras.layers import Dense
+from tensorflow.keras.models import Sequential
+
+import mlflow
+import mlflow.spark
+
+
+@click.command()
+@click.option("--ratings-data", help="Path readable by Spark to the ratings Parquet file")
+@click.option("--als-model-uri", help="Path readable by load_model to ALS MLmodel")
+@click.option("--hidden-units", default=20, type=int)
+def train_keras(ratings_data, als_model_uri, hidden_units):
+    np.random.seed(0)
+    tf.set_random_seed(42)  # For reproducibility
+
+    with pyspark.sql.SparkSession.builder.getOrCreate() as spark:
+        als_model = mlflow.spark.load_model(als_model_uri).stages[0]
+        ratings_df = spark.read.parquet(ratings_data)
+        (training_df, test_df) = ratings_df.randomSplit([0.8, 0.2], seed=42)
+        training_df.cache()
+        test_df.cache()
+
+        mlflow.log_metric("training_nrows", training_df.count())
+        mlflow.log_metric("test_nrows", test_df.count())
+
+        print(f"Training: {training_df.count()}, test: {test_df.count()}")
+
+        user_factors = als_model.userFactors.selectExpr("id as userId", "features as uFeatures")
+        item_factors = als_model.itemFactors.selectExpr("id as movieId", "features as iFeatures")
+        joined_train_df = training_df.join(item_factors, on="movieId").join(
+            user_factors, on="userId"
+        )
+        joined_test_df = test_df.join(item_factors, on="movieId").join(user_factors, on="userId")
+
+        # We'll combine the movies and ratings vectors into a single vector of length 24.
+        # We will then explode this features vector into a set of columns.
+        def concat_arrays(*args):
+            return list(chain(*args))
+
+        concat_arrays_udf = udf(concat_arrays, ArrayType(FloatType()))
+
+        concat_train_df = joined_train_df.select(
+            "userId",
+            "movieId",
+            concat_arrays_udf(col("iFeatures"), col("uFeatures")).alias("features"),
+            col("rating").cast("float"),
+        )
+        concat_test_df = joined_test_df.select(
+            "userId",
+            "movieId",
+            concat_arrays_udf(col("iFeatures"), col("uFeatures")).alias("features"),
+            col("rating").cast("float"),
+        )
+
+        pandas_df = concat_train_df.toPandas()
+        pandas_test_df = concat_test_df.toPandas()
+
+        # This syntax will create a new DataFrame where elements of the 'features' vector
+        # are each in their own column. This is what we'll train our neural network on.
+        x_test = pd.DataFrame(pandas_test_df.features.values.tolist(), index=pandas_test_df.index)
+        x_train = pd.DataFrame(pandas_df.features.values.tolist(), index=pandas_df.index)
+
+        # Show matrix for example.
+        print("Training matrix:")
+        print(x_train)
+
+        # Create our Keras model with two fully connected hidden layers.
+        model = Sequential()
+        model.add(Dense(30, input_dim=24, activation="relu"))
+        model.add(Dense(hidden_units, activation="relu"))
+        model.add(Dense(1, activation="linear"))
+
+        model.compile(loss="mse", optimizer=keras.optimizers.Adam(lr=0.0001))
+
+        early_stopping = EarlyStopping(
+            monitor="val_loss", min_delta=0.0001, patience=2, mode="auto"
+        )
+
+        model.fit(
+            x_train,
+            pandas_df["rating"],
+            validation_split=0.2,
+            verbose=2,
+            epochs=3,
+            batch_size=128,
+            shuffle=False,
+            callbacks=[early_stopping],
+        )
+
+        train_mse = model.evaluate(x_train, pandas_df["rating"], verbose=2)
+        test_mse = model.evaluate(x_test, pandas_test_df["rating"], verbose=2)
+        mlflow.log_metric("test_mse", test_mse)
+        mlflow.log_metric("train_mse", train_mse)
+
+        print(f"The model had a MSE on the test set of {test_mse}")
+        mlflow.tensorflow.log_model(model, name="keras-model")
+
+
+if __name__ == "__main__":
+    train_keras()