Airflow, MLflow or Kubeflow for MLOps?

Machine learning is now focusing more on the data (data-centric AI), and MLOps is obviously the way to bring ML projects into production. Today, I will review three popular MLOps frameworks, which are Airflow, MLflow and Kubeflow. The target is to answer two questions: what are the differences between them, and which frameworks should be used in modern MLOps systems.

I. Overview

At first, let's see into each framework to understand the features they can offer for your MLOps system.

1. Airflow

Airflow is a comprehensive framework for workflow orchestration (providing abilities to author, schedule and monitor workflows). A workflow can be your Directed Acyclic Graph (DAG) of your machine learning steps as below.

Machine learning pipeline/DAG

Machine learning pipeline/DAG

Airflow can be installed as a Python package, docker containers (docker-compose), or on a Kubenetes cluster. Airflow can be a good choice for both small and large scale systems. It is very Kubenetes-friendly: prividing well and easy integration with Kubenetes.

2. MLFlow

MLFlow is a tool for ML project management, experiment tracking, model management and registry. mlflow Python package provides API for managing ML experiments, runs, logging for parameters, metrics, artifacts and registering ML models. It also provides mechanism for version control of artifacts and models. This framework can also be installed as a Python package or in docker contaners (see installation). Mlflow can run in multiple scenerios, from local disk database, to remote tracking server, remote database server and separated storage host. Below are two examples of deployment scenerios of MLflow.

MLflow is deployed locally with file storage for both artifacts and database

MLflow is deployed locally with file storage for both artifacts and database

MLflow is deployed with separated tracking server, artifact host and remote database server

MLflow is deployed with separated tracking server, artifact host and remote database server

3. Kubeflow

While Airflow is a general workflow orchestration framework with no specific support for machine learning, and MLflow is a ML project management and tracking framework without a workflow orchestration system, Kubeflow is designed as a cloud-native platform that support all features for building MLOps: pipelines (workflow orchestration), training management and deployment. This platform is designed for large scale projects from the beginning and requires Kubenetes for operation.

4. Popularity

The popularity of a framework often tells us about the maturity and stalibility of that framework. A big community also brings rich documentation and support when we employ an open source in our project. We can see into the number of Github stars to estimate the popularity of an open source project and how cool it is. Below is the Github star history of Airflow, MLflow and Kubeflow.

Star History of Airflow, MLflow and Kubeflow at 26/03/2022 - Source:

Star History of Airflow, MLflow and Kubeflow at 26/03/2022 - Source: https://star-history.com/#kubeflow/kubeflow&mlflow/mlflow&apache/airflow&Date

Among three frameworks, Airflow has more Github stars (~25k at 26/03/2022), while Kubeflow and MLflow have half of that number (~12k). That can tells us that Airflow may be employed by more companies / projects tham MLflow and Kubeflow. The popularity of Airflow may because Airflow is used in many other tasks which need workflow orchestration (instead targeting only ML like Kubeflow or MLFlow).

II. Airflow + MLflow vs. Kubeflow

For full features of a MLOps system, Airflow needs to be combined with MLflow, while Kubeflow can almost provide all the features needed for a MLOps system. In this comparison, I also want to join Airflow with MLflow to build a MLOps stack. The other is Kubeflow. The figure beflow describe the features available in these stacks.

Airflow, MLflow and Kubeflow features

Airflow, MLflow and Kubeflow features

1. Workflow orchestration and data passing

Both Airflow and Kubeflow pipelines are workflow orchestration frameworks. However, Kubeflow pipelines is designed to support ML project better. For examples, it defines a standard way for passing data between machine learning operators using inputValue/outputValue or Kubeflow artifacts. In Airflow, we have XCOM as the key-value passing mechanism, however, we need to design and implement artifact passing ourself. Fortunately, when combining Airflow with MLflow, we can leverage artifact logging feature of MLflow for passing data between Airflow operators.

2. Experiment tracking and logging

We have experiment tracking, metric/artiface/model logging in MLflow and Kubeflow. However, the logging methods in mlfow packages seem to be easier to use than in Kubeflow. MLflow also provides ability to log and retrieve metric history (loss, accuracy, ...) during training, while I haven't found anything similar in Kubeflow.

3. Model registry and serving

MLflow has a great mechanism to register model easily by their name, while Kubeflow only support a complicated way for model register using ML metadata. However, model serving is supported better in Kubeflow with KServe or other addons.

3. Other ML specific features of Kubeflow

Kubeflow also provide us other features which we cannot find in Airflow + MLflow stacks, some of them are:

  • Kubeflow Notebooks, where user can create and run JupyterLab in their Kubenetes cluster, instead of local machine.
  • Katib, where user can run hyper parameter tuning or network architecture search at large scale.
  • External addons to deal with multiple data science and platform problems.

4. Scalibility

Airflow + MLflow stack is very flexible with running environment, from Python packages, native docker on your machine to a Kubenetes cluster/cloud. In contrast, Kubeflow always needs Kubenetes for up and running. That may because Kubeflow is design at large scale in mind. For this reason, I think Airflow + MLflow stack fits better for small scale systems, where we don't want to setup and maintain Kubenetes clusters. Here we can write anything in docker/docker-compose and deploy with just one command. Airflow also can be scaled for Kubenetes cloud by using KubernetesPodOperator or Kubenetes Executor. In contrast, Kubeflow needs Kubenetes (on premise or managed cloud) to setup and run. In exchange, you will have a stable system with full features for machine learning.

Airflow with Kubenetes Executor - Source:

Airflow with Kubenetes Executor - Source: https://airflow.apache.org/docs/apache-airflow/stable/executor/kubernetes.html

5. How to choose between Airflow+Mlflow, and Kubeflow?

To sum up, I have some recommendations from my personal perspective:

  • If your system needs to deal with multiple types of workflow, not just machine learning, Airflow may support you better. It is a mature workflow orchestration frameworks with support for a lot of operators besides machine learning.
  • If you want a system with predesigned patterns for machine learning, and run at large scale on Kubenetes clusters, you may want to consider Kubeflow. Many ML specific components in Kubeflow can save your time implementing from scratch in Airflow.
  • If you want to deploy MLOps in a small scale system (for example, a workstation, or a laptop), picking Airflow+MLflow stack can eliminate the need of setting up and running a Kubenetes system, and save more resources for the main tasks.

This blog post has briefly shown the differences between three popular MLOps frameworks (Airflow, MLflow and Kubeflow). Hope that it helps you in making decision between 2 stacks (Airflow + MLflow and Kubeflow). If you want to talk more about these frameworks or recommend others, please comment beflow. Thank you very much!

Related Posts

With Daisykit – Everyone can build AI projects!

Imagine you are a software engineer or a DIY hobbyist with great ideas to build AI-powered projects. However, it will be such difficulty when you know very little or nothing about AI.

Read more

Install ROS 2 on Raspberry Pi 4 (SD card image available)

ROS2 is an excellent framework for robotics applications. You can go further and embed a completed ROS 2 application stack in a small robot base with a Raspberry Pi computer.

Read more

Paper review: "YOLOX: Exceeding YOLO Series in 2021" and application in traffic sign detection - VIA Autonomous

YOLOX is an anchor-free version of YOLO, with a simpler design but better performance! It aims to bridge the gap between research and industrial communities.

Read more