👋 Hi! I’m Bibin Wilson. In each edition, I share practical tips, guides, and the latest trends in DevOps and MLOps to make your day-to-day DevOps tasks more efficient. If someone forwarded this email to you, you can subscribe here to never miss out!

✉️ In Today’s MLOps Edition

In the last edition, we learned how Kubeflow Trainer helps run distributed training workloads on Kubernetes.

Today, we will look at how MLflow helps track experiments, store models, and manage the model lifecycle using practical examples.

It covers,

  • What is Experiment Tracking?

  • Understanding the MLflow Architecture

  • How Kubeflow Pipelines Talk to MLflow

  • Deploying MLflow on Kubernetes (Hands-on)

  • Training an Employee Attrition Model Locally and Tracking Runs with MLflow (Hands-on)

  • Registering and Managing Models in the MLflow Model Registry

and more..

Planning to take a Kubernetes certification?

The Linux Foundation is currently running a limited-time sale.

Use code JULY26CCCT at kube.promo/devops to get flat 35% OFF CKA, CKAD, CKS, and other certifications.

For 45% bundle discounts and other offers, check this GitHub repository.

⏳ These offers are available for a limited time, so don't miss out.

Kubeflow Trainer takes care of creating training jobs, managing workers, scheduling resources, and running the training process.

But once model training is completed, we have a new problem.

How do we know:

  • Where do we store the trained model?

  • Which dataset created this model?

  • How to compare models from different training runs?

  • Which model version should go to production? and more..

This is where MLflow comes in.

Before we dive deep into MLflow, you need to understand what experiment tracking is. This way you will be able to relate to MLflow better.

Before You Continue: This is part of an ongoing MLOps series. You can check this repo to go through all the previous editions in order.

What is Experiment Tracking?

In ML, every training run is called an experiment. As I explained in the training edition, for training we pick an algorithm, set the hyperparameters, point it to a dataset version, and train. If something changes, we train again.

Experiment tracking means automatically logging important information about a training run. For example, during every training run, the following four things get captured.

  • Parameters: Algorithm, hyperparameters, dataset version.

  • Metrics: Accuracy, F1 score, training time.

  • Artifacts: The model file, environment files, etc

  • Metadata: Info on artifact location, model signature, input example, custom metadata, etc.

It is similar to how we track commit SHA, build number, logs, etc in the CI system. So a training run without tracking is like a CI build without build history.

The most widely used tool for experiment tracking is MLflow.

What is MLflow?

MLflow is the Git for machine learning experiments.

It is an open source platform used to track and manage the complete lifecycle of a model. Meaning, who trained it, with what parameters, on what data, what the results were, and which version is approved for production.

By tracking every training run, MLflow makes it easy for data scientists to compare experiments and identify which combination of data, parameters, and code produced the best-performing model.

Since all the metadata is stored, you can also reproduce a training run using the same code, parameters, and dataset.

Important Note: MLFlow has lot of LLM and agent related features (tracing, prompt registry, evaluation). Which I will cover in the future LLMOps editions. In this edition I am covering only the classic ML part.

The following diagram illustrates the components of MLflow.

MLflow operates as a client-server system, and the following are its four core components.

  1. MLflow SDK (Client): A Python package used to connect to the MLflow tracking server. You can install it on your local system or call it in ML training workflows to push the details to MLflow.

  2. Tracking server: A lightweight FastAPI-based web server with a UI and a REST API. Your training code sends data to it over HTTP.

  3. Backend store: A relational database that stores the metadata of experiments, runs, traces, etc. MLFlow supports PostgreSQL, MySQL, SQLite, and MSSQL as backends.

  4. Artifact store: This component is used to store artifacts such as model weights, images, data files, etc. For our attrition model project, this is where model.pkl ends up after training. The artifact store is implemented using object storage services such as AWS S3, Minio, GCS, Azure Blob Storage, etc.

MLFow Functional components

MLflow has the following four key functional components.

  1. MLflow Experiments: An experiment is a logical container for your ML work. For example, everything related to the attrition model falls under a single experiment.

  2. MLflow Run: A run is a single execution of your training code inside an experiment. Every run records the parameters, metrics, artifacts, and code version used for that training. If you train the model 50 times with different hyperparameters, you get 50 runs under one experiment.

  3. Model Registry: A centralized repository for managing versioned machine learning models. Once you identify the best-performing model from your experiment runs, you can register it in the Model Registry.

You will understand all these better in the hands-on covered in the upcoming sections.

How the Kubeflow Pipeline Talks to MLflow

Here is the common question everyone has.

We already have Airflow, DVC, Feast, and Kubeflow in our stack. Where exactly does MLflow fit in the MLOps stack?

The integration happens inside the training script that runs as part of a Kubeflow Pipeline.

When the training component executes, the script connects to the MLflow Tracking Server using the configured tracking URI.

During training, it logs all model-related information to MLflow, including hyperparameters, evaluation metrics, model artifacts, and metadata. Model artifacts are then pushed to the configured artifact store (such as Amazon S3).

MLflow on Kubernetes (EKS Hands-on)

Now let’s put everything we learned into practice with a hands-on setup. Here is what we will do.

  • We will set up MLflow on Kubernetes.

  • Configure AWS S3 as an artifact store.

  • Run our employee attrition training locally to train and log employee attrition model details.

  • Explore the experiment, run details, artifacts, and registered model in the MLflow UI.

Important Note: This setup guide is based on EKS cluster with S3 pod identity integration. If you are trying this our on different setup, you need to modify the S3 integration accordingly.

Follow the steps to set up MLflow on an EKS cluster.

Step 1: Create an S3 bucket

We will use an AWS S3 bucket as the backend Artifact store. First, you need to create a bucket. Replace dcube-mlflow-artifact-store in the following command with a unique bucket name.

aws s3api create-bucket \
  --bucket dcube-mlflow-artifact-store \
  --region us-west-2 \
  --create-bucket-configuration LocationConstraint=us-west-2 \
  --no-cli-pager

Verify the bucket exists.

aws s3api head-bucket \
  --bucket dcube-mlflow-artifact-store \
  --no-cli-pager

Step 2: Deploy PostgreSQL Using Helm Chart

Since we will be using PostgreSQL as the MLflow backend store, let's deploy it using Helm.

helm install mlflow-postgres oci://registry-1.docker.io/bitnamicharts/postgresql \
  --namespace mlflow \
  --create-namespace \
  --set auth.username=mlflow \
  --set auth.password=mlflow123 \
  --set auth.database=mlflow \
  --set primary.persistence.size=10Gi

Warning: Here I have given the username and password are given directly. In production, store the secrets in external secret managers like AWS SecretsManager and call it during deployment.

Ensure the Postgres pod is running.

$ kubectl get po -n mlflow

NAME                           READY   STATUS    RESTARTS   AGE
mlflow-postgres-postgresql-0   1/1     Running   0          35m

Step 3: Update the S3 bucket in the Helm Values File

The next step is to customize the Helm values for the MLflow installation.

You will find the mlflow.yaml values file inside the platform-tools/mlflow/helm folder.

In that file, replace <bucket-name> with your S3 bucket name in artifactsDestination.

The mlflow.yaml also contains the predefined PostgreSQL service endpoint we set up in the previous step. Also, we are enabling NodePort in the values to access the MLflow UI.

Step 4: Deploy MLflow

Now, let's deploy MLflow on the Kubernetes cluster using the Helm Chart.

helm install mlflow . --namespace mlflow -f mlflow.yaml

Run the following command to check if the pods are up and running.

$ kubectl get po -n mlflow

NAME                             READY   STATUS    RESTARTS   AGE
mlflow-mlflow-68cddb7f64-r277z   1/1     Running   0          98s
mlflow-postgres-postgresql-0     1/1     Running   0          58m

Step 5: Access the MLflow UI

You can access the MLflow UI using the NodePort we exposed. You can get the NodePort using the following command.

kubectl get svc mlflow-mlflow -n mlflow

Or you can use the following port-forward command. Keep this

kubectl port-forward deployment/mlflow-mlflow 5000:5000 -n mlflow

You should be able to access the UI at localhost:5000 in your browser, as shown below.

Step 6: Configure EKS Pod Identity for S3 Access

Next, you need to set up AWS EKS Pod Identity so that your MLflow deployment can securely access an S3 bucket without storing AWS credentials in Kubernetes.

All the steps required for the s3 role creation and eks pod identity association are part of the eks-s3.sh shell script in the platform-tools/mlflow folder.

Open the shell script and update the following variables shown in the image with your EKS cluster and bucket names.

Once updated, execute the shell script with create input.

$ chmod +x eks-s3.sh 

$ ./eks-s3.sh create

You should see all the roles and association details in the script output. Now let’s move on to the MLFlow components setup

Run the Training Script

Important: The training script used in this edition is only meant to help you understand how MLflow works.

In our final MLOps project, I will show you how to use MLflow with the employee attrition model, including logging parameters, metrics, and registering the trained model.

You will find a train_and_log_model.py script inside the phase-2-enterprise-setup/mlflow folder. This script does the following.

  • Trains a scikit-learn Gradient Boosting model using the employee attrition dataset.

  • Logs the model, training parameters, and evaluation metrics (Accuracy, Precision, Recall, F1-score, and ROC-AUC) to the MLflow Tracking Server.

  • Registers the trained model in the MLflow Model Registry.

Now, open the Python script and replace the MLFLOW_TRACKING_URI with the EKS node IP and NodePort address. Alternatively, if you're using kubectl port-forward, set it to http://127.0.0.1:5000 instead.

Now, create a Python environment and install the dependencies.

python3 -m venv venv                                    
source venv/bin/activate
pip install -r requirements.txt

Once installed, run the train_and_log_model Python script.

python train_and_log_model.py

Once you execute the script, the log_params(), log_metrics(), and set_tags() send metadata, while mlflow.sklearn.log_model() sends the actual trained model (model.skops) to MLFlow along with the files needed to reproduce and serve it.

Explore the Run in the UI

Now, if you visit the MLFlow UI, you will find the employee-attrition experiment created as shown below. Each execution of the training script creates a new run under this experiment.

If you click the employee-attrition experiment, it will show all the training runs. If you select a specific run, you can view everything logged during training, including Parameters, Metrics, Model artifacts, etc. The following animated GIF shows it better.

💡 Key Insight:

Every training run creates a model artifct and you can view them under models.

However, not all models become part of the Model Registry. Data scientists select the best model and register it so that it becomes part of the Model Registry. This process is covered in the section below.

If you expand the Artifacts section of a specific run and select the MLmodel metadata file, you can view the s3 artifact location.

If you browse the configured S3 bucket, you will find the same model files uploaded by MLflow during training.

And, if you scroll a little further, the MLmodel metadata file also contains the model signature, as shown below.

A model signature is the schema of a machine learning model. It defines the expected input features, their data types, the prediction output format, and optional inference parameters

💡 Key Insight:

In the feature store edition, I explained why feature order matters during inference. The signature solves this problem. It gives you the exact order of features that need to be passed to the model during inference.

Comparing Training Runs

Training is usually not a one-time activity. Data scientists train the model multiple times with different algorithms, hyperparameters, different datasets, etc.

Then they can select the runs in the UI and use the compare feature to determine which training run produced the best model, as shown below.

This makes it easy for data scientists to determine which run to deploy.

Register & Promote the Best Model in Model Registry

After comparing multiple runs, the data scientist team selects the best-performing model and registers it in the Model Registry.

This creates a new version of the registered model that stores its artifacts, metadata, metrics, lineage, and version history.

Let’s register a specific run using the register_model script.

python register_model.py

This script registers the model and assigns the production alias champion tag to the model.

The @champion alias acts like a :stable container image tag. It always points to the model version approved for production deployment. This allows deployment systems to reference models:/my-model@champion without hardcoding version numbers.

Key Insight:

In production MLOps setups, this entire workflow is automated through CI/CD pipelines.

After a model passes predefined quality gates such as accuracy ≥ 95%, F1 score ≥ 94%, and data quality validation, the pipeline automatically registers the model, promotes it with an alias like Champion, and deploys it to the inference platform.

Cleanup

Once you have tested the setup, run the following cleanup commands to clean up the setup.

$ helm uninstall mlflow -n mlflow

$ helm uninstall mlflow-postgres -n mlflow

$ ./eks-s3.sh delete

$ aws s3 rb s3://dcube-mlflow-artifact-store --force

What’s Coming Next?

Our best model in the registry is employee-attrition@champion. How do we actually serve it to real users at scale?

In the next edition, I will cover KServe. We will look at what KServe is, and how it pulls the registered @champion model from MLflow and deploys it on Kubernetes as a scalable inference service

Let’s continue building this step by step.

Reply

Avatar

or to participate

Keep Reading