MLflow: Experiment Tracking & Model Registry

Kubeflow Trainer takes care of creating training jobs, managing workers, scheduling resources, and running the training process. But once training finishes, a new problem shows up: where do we store the trained model, which dataset produced it, how do we compare models from different runs, and which version should go to production? That's where MLflow comes in.

What is Experiment Tracking?

Each time you train a model, a new training run is created. An experiment is a collection of those runs. As covered in the model training guide, you pick an algorithm, set hyperparameters, point it to a dataset version, and train — if something changes, you train again.

Experiment tracking means automatically logging four things about every training run:

• Parameters — algorithm, hyperparameters, dataset version
• Metrics — accuracy, F1 score, training time
• Artifacts — the model file, environment files, etc.
• Metadata — artifact location, model signature, input example, custom metadata

What is MLflow?

MLflow is the git for machine learning experiments — an open-source platform that tracks and manages the complete lifecycle of a model: who trained it, with what parameters, on what data, what the results were, and which version is approved for production. By tracking every run, it makes it easy to compare experiments and identify which combination of data, parameters, and code produced the best-performing model. Since the metadata is stored, you can also reproduce any run using the same code, parameters, and dataset.

MLflow operates as a client-server system with four core components:

MLflow Functional Components

1. Experiments — a logical container for your ML work. Everything related to the attrition model falls under a single experiment.
2. Runs — a single execution of your training code inside an experiment. Every run records the parameters, metrics, artifacts, and code version used. Train 50 times with different hyperparameters and you get 50 runs under one experiment.
3. Model Registry — a centralized repository for managing versioned models. Once you identify the best-performing run, you register its model in the Model Registry.

How Kubeflow Pipelines Talk to MLflow

With Airflow, DVC, Feast, and Kubeflow already in the stack, where does MLflow fit? 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, and during training it logs hyperparameters, evaluation metrics, model artifacts, and metadata — with model artifacts pushed to the configured artifact store (such as S3).

Hands-on: MLflow on Kubernetes (EKS)

We'll set up MLflow on Kubernetes, configure S3 as the artifact store, run a training script locally, and explore the results in the MLflow UI.

1. Create an S3 bucket

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

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

2. Deploy PostgreSQL as the backend store

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

Confirm the Postgres pod is running:

3. Point the Helm values at your S3 bucket

In the MLflow Helm values file, set artifactsDestination to your S3 bucket name, and enable a NodePort so the UI is reachable.

4. Deploy MLflow

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

$ 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

5. Access the MLflow UI

# Via NodePort
kubectl get svc mlflow-mlflow -n mlflow

# Or via port-forward
kubectl port-forward deployment/mlflow-mlflow 5000:5000 -n mlflow

The UI is then reachable at localhost:5000.

6. Configure EKS Pod Identity for S3 access

Set up EKS Pod Identity so the MLflow deployment can access the S3 bucket securely without storing AWS credentials in Kubernetes. Update the cluster and bucket names in the setup script, then run it:

chmod +x eks-s3.sh
./eks-s3.sh create

7. Run the training script

Update MLFLOW_TRACKING_URI to the EKS node IP + NodePort address (or http://127.0.0.1:5000 if using kubectl port-forward), then train:

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

python train_and_log_model.py

This trains a scikit-learn gradient boosting model on the employee attrition dataset, logs parameters, metrics (accuracy, precision, recall, F1, ROC-AUC), and the model to the tracking server, and registers the trained model in the Model Registry. log_params(), log_metrics(), and set_tags() send metadata, while mlflow.sklearn.log_model() sends the trained model artifact along with the files needed to reproduce and serve it.

8. Explore the run in the UI

The MLflow UI shows an employee-attrition experiment; each execution of the training script creates a new run under it. Selecting a run shows everything logged during training — parameters, metrics, and model artifacts.

Model Signature

Every model expects inputs in a certain schema and format. If you pass inputs that don't match, you can get wrong predictions with no obvious error. A model signature is the schema of a model — the expected input features, their data types, the prediction output format, and optional inference parameters. It's stored in the run's MLmodel metadata file alongside the S3 artifact location.

For any consumer, the signature acts as a contract: it tells you exactly what the model expects as input and what output format to expect — the model equivalent of an API spec. CI/CD tooling can use it to compare schema changes against the production model before deployment.

Register & Promote a Champion Model

Training isn't a one-time activity — teams train with different algorithms, hyperparameters, and datasets, then use the compare feature in the UI to determine which run produced the best model.

After comparing runs, register the best-performing one:

python register_model.py

This creates a new version of the registered model with its artifacts, metadata, metrics, lineage, and version history, and assigns it the champion alias.

Clean up

helm uninstall mlflow -n mlflow
helm uninstall mlflow-postgres -n mlflow
./eks-s3.sh delete

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