• Key Features
• Getting Started
  • Installation
  • Model Development
  • Productionization
• Reference Guide

The MLflow Regression Recipe is an MLflow Recipe (previously known as MLflow Pipeline) for developing high-quality regression models. It is designed for developing models using scikit-learn and frameworks that integrate with scikit-learn, such as the XGBRegressor API from XGBoost.

This repository is a template for developing production-ready regression models with the MLflow Regression Recipe. It provides a recipe structure for creating models as well as pointers to configurations and code files that should be filled in to produce a working recipe.

Code developed with this template should be run with MLflow Recipes. An example implementation of this template can be found in the Recipe Regression Example repo, which targets the NYC taxi dataset for its training problem.

Note: MLflow Recipes is an experimental feature in MLflow. If you observe any issues, please report them here. For suggestions on improvements, please file a discussion topic here. Your contribution to MLflow Recipes is greatly appreciated by the community!

• Deterministic data splitting
• Reproducible data transformations
• Hyperparameter tuning support
• Model registration for use in production
• Starter code for ingest, split, transform and train steps
• Cards containing step results, including dataset profiles, model leaderboard, performance plots and more

To use this MLflow regression recipe, simply install the packages listed in the requirements.txt file. Note that Python 3.8 or above is recommended.

pip install requirements.txt

You may need to install additional libraries for extra features:

• Hyperopt is required for hyperparameter tuning.
• PySpark is required for distributed training or to ingest Spark tables.
• Delta is required to ingest Delta tables. These libraries are available natively in the Databricks Runtime for Machine Learning.

After installing MLflow Recipes, you can clone this repository to get started. Simply fill in the required values annotated by FIXME::REQUIRED comments in the Recipe configuration file and in the appropriate profile configuration: local.yaml (if running locally) or databricks.yaml (if running on Databricks).

The Recipe will then be in a runnable state, and when run completely, will produce a trained model ready for batch scoring, along with cards containing detailed information about the results of each step. The model will also be registered to the MLflow Model Registry if it meets registration thresholds. To iterate and improve your model, follow the MLflow Recipes usage guide. Note that iteration will likely involve filling in the optional FIXMEs in the step code files with your own code, in addition to the configuration keys.

Once the model under development has reached desired quality, one may programatically run the adapted recipe to either reproduce the model consistently, or train a production model with different input data source (e.g. production dataset) via automated processes such as continuous deployment (CD). Below are recommendations for productionizing a recipe depending on its running environments:

• [Local|CICD] Running the recipe via command line: mlflow recipes run --profile {PROFILE}
• [Databricks] Creating a new notebook (e.g. notebooks/databricks_prod.py), instantiate and run the entire recipe via r = Recipe(profile=PROFILE); r.run() with desired profile.

Below is a visual overview of the MLflow Regression Recipe's information flow.

Model develompent consists of the following sequential steps:

ingest -> split -> transform -> train -> evaluate -> register

The batch scoring workflow consists of the following sequential steps:

ingest_scoring -> predict

A detailed reference for each step follows.

• Reference
  • Step artifacts
  • Ingest step
  • Split step
  • Transform step
  • Train step
  • Evaluate step
  • Register step
  • Batch scoring
    • Ingest Scoring step
    • Predict step
  • MLflow Tracking
  • Model Registry configuration
  • Metrics
    • Built-in metrics
    • Custom metrics

Each of the steps in the recipe produces artifacts after completion. These artifacts consist of cards containing detailed execution information, as well as other step-specific information. The Recipe.inspect() API is used to view step cards. The get_artifact API is used to load all other step artifacts by name.
Per-step artifacts are further detailed in the following step references.

Specifies the name of the column containing targets / labels for model training and evaluation

• target_col: string. Required.
  The name of the column set as the target for model training. Example:

  target_col: "fare_amount"
  

The primary evaluation metric is the one that will be used to select the best performing model in the MLflow UI as well as in the train and evaluation steps. This can be either a built-in metric or a custom metric (see below).
Models are ranked by this primary metric.

• primary_metric: string. Required.
  The name of the primary evaluation metric. Example:

  primary_metric: "root_mean_squared_error"
  

The ingest step loads the dataset specified by the steps.ingest section in recipe.yaml.
Note: If you make changes to the dataset referenced by the ingest step (e.g. by adding new records or columns), you must manually re-run the ingest step in order to use the updated dataset in the recipe. The ingest step does not automatically detect changes in the dataset.

Below are all the possible options and full reference guide for different configurations allowed in the ingest step:

steps:
  ingest:
    using: "parquet"
    location: "./data/sample.parquet"

steps:
  ingest:
    using: "delta"
    location: "dbfs:/user/hive/warehouse/ml.db/all_prices_1990_2021_mar"
    version: 1

steps:
  ingest:
    using: "spark_sql"
    sql: "SELECT * FROM delta.`dbfs:/databricks-datasets/nyctaxi-with-zipcodes/subsampled`"

steps:
  ingest:
    using: "custom"
    location: "./data/sample.csv"
    loader_method: read_csv_as_dataframe

def read_csv_as_dataframe(location: str) -> DataFrame:
  import pandas

  return pandas.read_csv(location, index_col=0)

Step artifacts:

• ingested_data: The ingested data as a Pandas DataFrame.

The split step splits the ingested dataset produced by the ingest step into:

• a training dataset for model training
• a validation dataset for model performance evaluation & tuning, and
• a test dataset for model performance evaluation.

Subsequent steps use these datasets to develop a model and measure its performance.

The post-split method should be written in steps/split.py and should accept three parameters:

• train_df: DataFrame. The unprocessed train dataset.
• validation_df: DataFrame. The unprocessed validation dataset.
• test_df: DataFrame. The unprocessed test dataset.

It should return a triple representing the processed train, validation and test datasets. steps/split.py contains an example placeholder function.

The split step is configured by the steps.split section in recipe.yaml as follows:

Step artifacts:

• training_data: the training dataset as a Pandas DataFrame.
• validation_data: the validation dataset as a Pandas DataFrame.
• test_data: the test dataset as a Pandas DataFrame.

The transform step uses the training dataset created by the split step to fit a transformer that performs the user-defined transformations. The transformer is then applied to the training dataset and the validation dataset, creating transformed datasets that are used by subsequent steps for estimator training and model performance evaluation.

The transform step is configured by the steps.transform section in recipe.yaml: Note: The steps.transform section is not required. If absent, an identity transformer will be used.

Below are all the possible options and full reference guide for different configurations allowed in the transform step:

Step artifacts:

• transformed_training_data: transformed training dataset as a Pandas DataFrame.
• transformed_validation_data: transformed validation dataset as a Pandas DataFrame.
• transformer: the sklearn transformer.

The train step uses the transformed training dataset output from the transform step to fit an AutoML or user-defined estimator. The estimator is then joined with the fitted transformer output from the transform step to create a model recipe. Finally, this model recipe is evaluated against the transformed training and validation datasets to compute performance metrics.

Custom evaluation metrics are computed according to definitions in steps/custom_metrics.py and the metrics section of recipe.yaml; see Custom Metrics section for reference. Hyper-parameter tuning during the train step can also be enabled as a step config in the steps.train section in recipe.yaml.

The model recipe and its associated parameters, performance metrics, and lineage information are logged to MLflow Tracking, producing an MLflow Run.

The train step is configured by the steps.train section in recipe.yaml.

Below are all the possible options and full reference guide for different configurations allowed in the train step:

steps:
  train:
    using: "automl/flaml"
    time_budget_secs: 3000

steps:
  train:
    using: "custom"
    estimator_method: estimator_fn
    estimator_params:
      depth: 3
    tuning:
      enabled: True
      algorithm: "hyperopt.rand.suggest"
      max_trials: 5
      early_stop_fn: my_early_stop_fn
      parallelism: 1
      sample_fraction: 0.5
      parameters:
          alpha:
            distribution: "uniform"
            low: 0.0
            high: 0.01
          penalty:
            values: ["l1", "l2", "elasticnet"]

from typing import Dict, Any

def estimator_fn(estimator_params: Dict[str, Any] = None):
    from sklearn.linear_model import SGDRegressor

    return SGDRegressor(random_state=42, **estimator_params)

def my_early_stop_fn(*args):
  from hyperopt.early_stop import no_progress_loss
  return no_progress_loss(10)(*args)

Step artifacts:

• model: the MLflow Model recipe created in the train step as a PyFuncModel instance.

The evaluate step evaluates the model recipe created by the train step on the test dataset output from the split step, computing performance metrics and model explanations.

Performance metrics are compared against configured thresholds to produce a model_validation_status, which indicates whether or not a model is validated to be registered to the MLflow Model Registry by the subsequent register step.
These model performance thresholds are defined in the validation_criteria section of the evaluate step definition in recipe.yaml. Custom evaluation metrics are computed according to definitions in steps/custom_metrics.py and the metrics section of recipe.yaml; see the custom metrics section for reference.

Model performance metrics and explanations are logged to the same MLflow Tracking Run used by the train step.

The evaluate step is configured by the steps.evaluate section in recipe.yaml:

Step artifacts:

• run: the MLflow Tracking Run containing the model recipe, as well as performance and metrics created during the train and evaluate steps.

The register step checks the model_validation_status output of the preceding evaluate step and, if model validation was successful (if model_validation_status is 'VALIDATED'), registers the model recipe created by the train step to the MLflow Model Registry. If the model_validation_status does not indicate that the model passed validation checks (if model_validation_status is 'REJECTED'), the model recipe is not registered to the MLflow Model Registry.
If the model recipe is registered to the MLflow Model Registry, a registered_model_version is produced containing the model name and the model version.

The register step is configured by the steps.register section in recipe.yaml:

Step artifacts:

• registered_model_version: the MLflow Model Registry ModelVersion registered in this step.

After model training, the regression recipe provides the capability to score data with the trained model. This is convenient for scoring a static dataset, e.g., the test set from a Kaggle competition. For more advanced ETL workflows for scoring, we recommend using spark_udf directly instead, because it is more flexible to handle different data scenarios.

The ingest scoring step, defined in the steps.ingest_scoring section in recipe.yaml, specifies the dataset used for batch scoring and has the same API as the ingest step.

Step artifacts:

• ingested_scoring_data: the ingested scoring data as a Pandas DataFrame.

The predict step uses the model registered by the register step to score the ingested dataset produced by the ingest scoring step and writes the resulting dataset to the specified output format and location. To fix a specific model for use in the predict step, provide its model URI as the model_uri attribute of the recipe.yaml predict step definition.

The predict step is configured by the steps.predict section in recipe.yaml:

Step artifacts:

• scored_data: the scored dataset, with model predictions under the prediction column, as a Pandas DataFrame.

The MLflow Tracking server can be configured to log MLflow runs to a specific server. Tracking information is specified in the profile configuration files - profiles/local.yaml if running locally and profiles/databricks.yaml if running on Databricks.

Configuring a tracking server is optional. If this configuration is absent, the default experiment will be used.

Tracking information is configured with the experiment section in the profile configuration:

experiment:
  name: "experiment-name"
  tracking_uri: "sqlite:///metadata/mlflow/mlruns.db"
  artifact_location: "./metadata/mlflow/mlartifacts"

To register trained models to the MLflow Model Registry, further configuration may be required. If unspecified, models will be logged to the same server as specified in the tracking URI.

To register models to a different server, specify the desired server in the model_registry section in the profile configuration:

model_registry:
  registry_uri: "sqlite:///metadata/mlflow/registry.db"
  model_name: "model-name"

Evaluation metrics calculate model performance against different datasets. The metrics defined in the recipe will be calculated as part of the training and evaluation steps, and calculated values will be recorded in each step’s information card.

This regression recipe features a set of built-in metrics, and supports user-defined metrics as well.

Metrics are configured under the custom_metrics section of recipe.yaml, according to the following specification:

Note that each metric specifies a boolean value greater_is_better, which indicates whether a higher value for that metric is associated with better model performance.

The following metrics are built-in. Note that greater_is_better = False for all these metrics:

• mean_absolute_error
• mean_squared_error
• root_mean_squared_error
• max_error
• mean_absolute_percentage_error

Custom evaluation metrics define how trained models should be evaluated against custom criteria not captured by built-in sklearn evaluation metrics.

Custom evaluation metric functions should be defined in steps/custom_metrics.py. Function to define the custom_metrics should adhere to the following signature.

Custom metrics are specified as a list under the metrics.custom key in recipe.yaml, specified as follows:

• name: string. Required.
  Name of the custom metric. This will be the name by which you refer to this metric when including it in model evaluation or model training.

• function: string. Required. Specifies the function this custom metric refers to.

• greater_is_better: boolean. Required. Boolean indicating whether a higher metric value indicates better model performance.

An example custom metric configuration is as follows:

custom_metrics:
 - name: weighted_mean_square_error
   function: get_custom_metrics
   greater_is_better: True