Chapter 8: Training Models with MLflow

In this chapter, you will learn about creating production-ready training jobs with MLflow. In the bigger scope of things, we will focus on how to move from the training jobs in the notebook environment that we looked at in the early chapters to a standardized format and blueprint to create training jobs.

Specifically, we will look at the following sections in this chapter:

• Creating your training project with MLflow
• Implementing the training job
• Evaluating the model
• Deploying the model in the Model Registry
• Creating a Docker image for your training job

It's time to add to the pyStock machine learning (ML) platform training infrastructure to take proof-of-concept models created in the workbench developed in Chapter 3, Your Data Science Workbench to a Production Environment.

In this chapter, you will be developing a training project that runs periodically or when triggered by a dataset arrival....

Technical requirements

For this chapter, you will need the following prerequisites:

• The latest version of Docker installed on your machine. If you don't already have it installed, please follow the instructions at https://docs.docker.com/get-docker/.
• The latest version of Docker Compose installed—please follow the instructions at https://docs.docker.com/compose/install/.
• Access to Git in the command line, and installed as described at https://git-scm.com/book/en/v2/Getting-Started-Installing-Git.
• Access to a Bash terminal (Linux or Windows).
• Access to a browser.
• Python 3.5+ installed.
• The latest version of your ML library installed locally as described in Chapter 4, Experiment Management in MLflow

Creating your training project with MLflow

You receive a specification from a data scientist based on the XGBoost model being ready to move from a proof-of-concept to a production phase.

We can review the original Jupyter notebook from which the model was registered initially by the data scientist, which is a starting point to start creating an ML engineering pipeline. After initial prototyping and training in the notebook, they are ready to move to production.

Some companies go directly to productionize the notebooks themselves and this is definitely a possibility, but it becomes impossible for the following reasons:

• It's hard to version notebooks.
• It's hard to unit-test the code.
• It's unreliable for long-running tests.

With these three distinct phases, we ensure reproducibility of the training data-generation process and visibility and clear separation of the different steps of the process.

We will start by organizing our MLflow project...

Implementing the training job

We will use the training data produced in the previous chapter. The assumption here is that an independent job populates the data pipeline in a specific folder. In the book's GitHub repository, you can look at the data in https://github.com/PacktPublishing/Machine-Learning-Engineering-with-MLflow/blob/master/Chapter08/psystock-training/data/training/data.csv.

We will now create a train_model.py file that will be responsible for loading the training data to fit and produce a model. Test predictions will be produced and persisted in the environment so that other steps of the workflow can use the data to evaluate the model.

The file produced in this section is available at the following link:

https://github.com/PacktPublishing/Machine-Learning-Engineering-with-MLflow/blob/master/Chapter08/psystock-training/train_model.py:

1. We will start by importing the relevant packages. In this case, we will need pandas to handle the data, xgboost...

Evaluating the model

We will now move on to collect evaluation metrics for our model, to add to the metadata of the model.

We will work on the evaluate_model.py file. You can follow along by working in an empty file or by going to https://github.com/PacktPublishing/Machine-Learning-Engineering-with-MLflow/blob/master/Chapter08/psystock-training/evaluate_model.py. Proceed as follows:

1. Import the relevant packages—pandas and mlflow—for reading and running the steps, respectively. We will rely on importing a selection of model-evaluation metrics available in sklearn for classification algorithms, as follows:

   import pandas as pd
   import mlflow
   from sklearn.model_selection import train_test_split
   from sklearn.metrics import  \
       classification_report, \
       confusion_matrix, \
       accuracy_score, \
       auc, \
       average_precision_score, \
      ...

Deploying the model in the Model Registry

Next, you should add the register_model.py function to register the model in the Model Registry.

This is as simple as executing the mlflow.register_model method with the Uniform Resource Identifier (URI) of the model and the name of the model. Basically, a model will be created if it doesn't already exist. If it's already in the registry, a new version will be added, allowing the deployment tools to look at the models and trace the training jobs and metrics. It also allows a decision to be made as to whether to promote the model to production or not. The code you'll need is illustrated in the following snippet:

import mlflow
if __name__ == "__main__":
    
    with mlflow.start_run(run_name="register_model") as run:
        mlflow.set_tag("mlflow.runName", "register_model")
     ...

Creating a Docker image for your training job

A Docker image is, in many contexts, the most critical deliverable of a model developer to a more specialized systems infrastructure team in production for a training job. The project is contained in the following folder of the repository: https://github.com/PacktPublishing/Machine-Learning-Engineering-with-MLflow/tree/master/Chapter08/psystock-training-docker. In the following steps, we will produce a ready-to-deploy Docker image of the code produced:

1. You need to set up a Docker file in the root folder of the project, as shown in the following code snippet:

   FROM continuumio/miniconda3:4.9.2
   RUN apt-get update && apt-get install build-essential -y
   RUN pip install \
       mlflow==1.18.0 \
       pymysql==1.0.2 \
       boto3
   COPY ./training_project /src
   WORKDIR /src

2. We will start by building and training the image by running the following command:

   docker build -t psystock_docker_training_image...

Summary

In this chapter, we introduced the concepts and different features in terms of using MLflow to create production training processes.

We started by setting up the basic blocks of the MLflow training project and followed along throughout the chapter to, in sequence, train a model, evaluate a trained model, and register a trained model. We also delved into the creation of a ready-to-use image for your training job.

This was an important component of the architecture, and it will allow us to build an end-to-end production system for our ML system in production. In the next chapter, we will deploy different components and illustrate the deployment process of models.

Further reading

In order to further your knowledge, you can consult the official documentation at the following link:

https://www.mlflow.org/docs/latest/projects.html