MLflow, Airflow, and Kubeflow

MLflow

MLflow is an open source platform for tracking experiments and storing models. It covers two needs from the pipeline lesson: run metadata and the model registry.

MLflow runs as a server with a web UI. It stores metadata in a database, for example Postgres. Model files live in a separate store, for example S3. Locally, the server starts with mlflow server.

Run: recording an experiment

A data scientist adds a few lines to the training code:

import mlflow
from sklearn.linear_model import LogisticRegression

mlflow.set_tracking_uri("http://mlflow:5000")  # MLflow server address
mlflow.set_experiment("churn")

with mlflow.start_run():                        # start a run
    model = LogisticRegression(C=0.5).fit(X_train, y_train)
    mlflow.log_param("C", 0.5)                  # training parameter
    mlflow.log_metric("accuracy", model.score(X_test, y_test))
    mlflow.sklearn.log_model(model, name="model", registered_model_name="churn")

Each start_run creates a run, one record of a training session. A run stores the parameters, metrics, and the model file. In the UI, runs show up as a table, and you can sort them by accuracy and compare them. Twenty attempts with different parameters no longer get lost in a notebook.

The model registry in MLflow

The registered_model_name="churn" argument puts the model into the model registry as a new version of the churn model. An alias, for example champion, is then assigned to a version.

The service loads the model by alias:

model = mlflow.pyfunc.load_model("models:/churn@champion")

Move the alias to another version, and after a restart the service loads the new model. MLflow can also spin up an HTTP service for the model on its own, which is handy for a quick check:

mlflow models serve -m "models:/churn@champion" -p 5001

Orchestrator

The training pipeline's steps need to run in order, on a schedule, and retry on failure. That's the job of an orchestrator. A training pipeline starts on a schedule or an event rather than on a commit, and its steps can run for hours. That's why it's usually run on a dedicated orchestrator rather than a CI system.

Airflow

Apache Airflow is a general-purpose open source task orchestrator. A pipeline in it is described in Python as a DAG (directed acyclic graph), a graph of tasks with no cycles. Here's what a DAG looks like in Airflow 3:

from airflow.sdk import dag, task
import pendulum

@dag(schedule="@weekly", start_date=pendulum.datetime(2026, 1, 1), catchup=False)
def retrain_churn():
    @task
    def extract() -> str: ...          # export the data, return the path
    @task
    def validate(path: str) -> str: ...
    @task
    def train(path: str) -> None: ...

    train(validate(extract()))         # order: extract → validate → train

retrain_churn()

Airflow runs this DAG once a week, shows each task's status in the UI, and retries failed ones if configured to. Besides ML, it's used for regular data exports too.

Kubeflow Pipelines

Kubeflow Pipelines is a pipeline orchestrator for Kubernetes. A pipeline step in it is called a component. Per the Kubeflow documentation, every component run executes as a separate container.

So each step has its own image with its own dependencies, and its own resources. The training step can get a GPU, while the data extraction step gets a bit of memory. The pipeline is described in Python through the Kubeflow Pipelines SDK.

Other open source orchestrators used for this kind of work: Argo Workflows, Prefect, Dagster.

What to remember

MLflow records each training session as a run: parameters, metrics, the model. The model registry in MLflow stores model versions, and the service loads it by alias, for example models:/churn@champion.

An orchestrator runs the pipeline's steps in order and on schedule. In Airflow, a pipeline is described as a Python DAG. In Kubeflow Pipelines, each step runs as a separate container in Kubernetes with its own resources.