Orchestrating machine learning experiments for MLOps using Apache Airflow
By: Andrea Capuano
Tags: Machine Learning, Airflow, Continuous Delivery, Mlops, Continuous Integration

Nowadays that more and more machine learning models are going to production, the need to operationalize the overall Machine Learning workflow becomes crucial to companies who adopt artificial intelligence capabilities.

We can use Apache Airflow platform to orchestrate the different phases of machine learning

Machine learning experiments usually follow a predefined set of phases, such as:

Data ingestion : Collect and integrate data from different sources

: Collect and integrate data from different sources Data validation : Ensure the collected data is valid and consistent with expectations

: Ensure the collected data is valid and consistent with expectations Data preparation : Validate, preprocess, extract features and transform the data to get it ready for the machine learning task

: Validate, preprocess, extract features and transform the data to get it ready for the machine learning task Model training : actual training of machine learning models, hyperparameters tuning

: actual training of machine learning models, hyperparameters tuning Model evaluation : Evaluate model performance, accept or reject its results

: Evaluate model performance, accept or reject its results Model deployment: Deploy models if the performances of the previous step are acceptable to go in production

Orchestrating the different phases through a well defined and repeatable workflow can boost up productivity in your overall machine learning pipelines. By both promoting well structured codebase and creating a way to reproduce systematically the steps. Hence, provide capabilities such as Continuous training and improvement of machine learning models. For example, we can automatically re-trigger the pipelines either at scheduled intervals or upon model performance degradation.

Meet Apache Airflow

Airflow is an open-source project backed up by the Apache software foundation. According to the official website, Airflow is:

A platform created by the community to programmatically author, schedule and monitor workflows.

Airflow is built with simplicity in mind, and allows to build directed acyclic graphs (DAG) completely in Python and by writing concise and understandable code. It is widely adopted in the industry for many different purposes and also is provided in some managed versions by some cloud providers (e.g: gcp cloud composer).

By using Airflow we can achieve ML pipeline orchestration quite easily, in the following paragraphs we’ll explore how to setup Airflow using Docker and write our first DAG.

Setting up Airflow with Docker and docker-compose

There are different ways to deploy Apache Airflow, I find Docker to be one of the easiest ones, for this article I will refer to this github repository, simply clone the repository and docker-compose up as shown below:



cd docker-airflow

docker-compose -f docker-compose-LocalExecutor.yml up git clone https://github.com/puckel/docker-airflow.git cd docker-airflowdocker-compose -f docker-compose-LocalExecutor.yml up

After a brief startup phase, Airflow will be up and running and you can visit the UI at the following address: http://localhost:8080, you will see something in the line of the following picture:

First look at the Airflow UI: We have one DAG called tutorial, turned off by default and having a specific schedule.

Creating our DAG

To create a Directed Acyclic Graph we need to provide a python file containing all the code to declare and instantiate the DAG. Let’s create our orchestrated_experiment.py file under the dags folder. Let’s then instantiate our DAG:

Now, we can create the different tasks that compose our DAG. The tasks will be exactly the ML phases we discussed early on in the article. To achieve this, we will need to use operators, they come in a big variety and they are used to create tasks. For example we can make use of the PythonOperator, which takes as input a python_callable: a python function that will be executed for that given task.

We’ll get the following result in Airflow:

ML steps inside Airflow, consisting of several PythonOperators together

Now, all you need to do is to define these functions and insert all the core logic of your experiment, using your favourite tools and frameworks. Such as pandas, scikit-learn, tensorflow and so on.

Examples of Orchestrated experiment.

How would we map ML tasks to operators? Let’s see some of the possibilities:

Data ingestion : We could use Hooks to connect to heterogeneus data sources, such as BigTable, Azure blob storage, ftp servers or your local filesystem

: We could use Hooks to connect to heterogeneus data sources, such as BigTable, Azure blob storage, ftp servers or your local filesystem Data validation: We can use a PythonOperator to ensure that the data is valid and consistent with what we expect. If we place assertions inside the functions and they fail, Airflow will mark automatically the task and consequently the DAG as failed

We can use a PythonOperator to ensure that the data is valid and consistent with what we expect. If we place inside the functions and they fail, Airflow will mark automatically the task and consequently the DAG as failed Data preparation: Use a PythonOperator to perform data preparation. You can install custom libraries such as pandas or scikit-learn by simply placing a requirements.txt file in the root directory of airflow and mapping a volume in the docker-compose volumes section, as follows: ./requirements.txt:/requirements.txt

Use a PythonOperator to perform data preparation. You can install custom libraries such as pandas or scikit-learn by simply placing a requirements.txt file in the root directory of airflow and mapping a volume in the docker-compose volumes section, as follows: ./requirements.txt:/requirements.txt Model training: We can still use a PythonOperator in combination with your favorite library (e.g: scikit-learn, tensorflow..), or you could get creative and use something like SageMaker training operator. If you want to keep your training separate and isolated, you can use a DockerOperator to run commands inside the training container and pass a reference to the location of the output of the data preparation step.

We can still use a PythonOperator in combination with your favorite library (e.g: scikit-learn, tensorflow..), or you could get creative and use something like SageMaker training operator. If you want to keep your training separate and isolated, you can use a to run commands inside the training container and pass a reference to the location of the output of the data preparation step. Model evaluation : This can be easily achieved by using the previously built model, that can be uploaded to a s3 bucket or a model registry (e.g: using MlFlow) and evaluate the performances on the test set.

: This can be easily achieved by using the previously built model, that can be uploaded to a s3 bucket or a model registry (e.g: using MlFlow) and evaluate the performances on the test set. Model deployment: You can start the deployment phase by checking the results of the evaluation step and if they are acceptable start the deployment. To achieve this you can use a PythonBranchOperator using the output of your favourite metric of choice (e.g: precision) on the value produced by the evaluation step.

You can easily check the outputs of previous steps by returning the values in the functions and retrieving them through XCOMs.

Conclusions

Suppose you have an algorithm running in production. You need to ensure that the precision of the algorithm doesn’t fall beneath 90% threshold. How can Airflow help?

We can have scheduled evaluations at regular times during the day and ensure that if the KPI requirement is not met an automatic retraining and redeployment is performed.

If companies do not operationalize these steps, the time-to-react to model degradation will highly depend on the speed of data scientists to perform (again) manual tasks that were already done in experimental phases.

Airflow is widely used and production ready, can scale out easily and fits well in many ecosystems thanks to the wide support of operators. However, make sure you also check out other interesting tools such as KubeFlow, PolyAxon or metaflow.