Track ML Experiments with MLflow

Machine learning projects involve many experiments. You change a learning rate. You swap a model. You tweak the batch size. Soon you cannot remember which run produced the best result.

MLflow solves this problem. It is an open-source platform that helps you track experiments, log parameters and metrics, and save models. This guide shows you how to use it from scratch.

What Is MLflow?

MLflow is a tool for managing the machine learning lifecycle. It has four main components. These are Tracking, Projects, Models, and Registry.

This article focuses on MLflow Tracking. It records everything about an experiment. That includes parameters, metrics, artifacts, and source code versions.

Each experiment is a named group of runs. Each run is a single execution of your training script. MLflow stores all run data in a local folder or a remote server.

Why Track ML Experiments?

Manual tracking does not scale. Spreadsheets get out of sync. Notebook outputs get lost. Team members repeat work by accident.

Proper tracking gives you clear benefits:

Reproducibility. You can recreate any run exactly. Every parameter and dataset version is recorded.

Comparison. You can compare hundreds of runs side by side. The best model becomes obvious.

Collaboration. A shared tracking server lets the whole team see results. No one works in isolation.

Auditability. You know who ran what, when, and with which settings.

Installing MLflow

Installation is simple. Use pip in your terminal.

pip install mlflow

Check the version to confirm the install.

mlflow --version

mlflow, version 2.11.3

If you also work with data pipelines, you may want to log pipeline outputs too. Tools like Kafka message producers and consumers in Python can feed live data into your experiments.

Starting the MLflow Tracking Server

MLflow runs a local tracking server by default. You can start the UI with one command.

mlflow ui

Open your browser at http://127.0.0.1:5000. You will see an empty dashboard. It fills up as you log runs.

For team use, start a remote server with a database backend.

mlflow server --backend-store-uri sqlite:///mlflow.db --default-artifact-root ./artifacts --host 0.0.0.0

Your First Tracked Run

Let us log a simple experiment. The code below trains a small model and records its results.

import mlflow
import mlflow.sklearn
from sklearn.linear_model import LogisticRegression
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score

# Load and split the data
data = load_iris()
X_train, X_test, y_train, y_test = train_test_split(
    data.data, data.target, test_size=0.2, random_state=42
)

# Start a run inside the default experiment
with mlflow.start_run():
    # Log parameters (inputs)
    mlflow.log_param("solver", "lbfgs")
    mlflow.log_param("max_iter", 200)

    # Train the model
    model = LogisticRegression(solver="lbfgs", max_iter=200)
    model.fit(X_train, y_train)

    # Evaluate and log the metric
    preds = model.predict(X_test)
    acc = accuracy_score(y_test, preds)
    mlflow.log_metric("accuracy", acc)

    # Save the model as an artifact
    mlflow.sklearn.log_model(model, "model")

    print(f"Accuracy: {acc}")

Accuracy: 1.0

Refresh the MLflow UI. You will see a new run. It shows the parameters, the accuracy metric, and the saved model.

Understanding the Core Concepts

The tracking API has a few key pieces. Learn them once and the rest is easy.

Experiment. A logical group of runs. Use mlflow.set_experiment() to choose or create one. If you skip it, runs go to the "Default" experiment.

Run. One execution. The mlflow.start_run() context manager opens a run and closes it cleanly when the block ends.

Parameters. Inputs of your run. Log them with mlflow.log_param(). Parameters are strings or numbers and cannot be changed after logging.

Metrics. Outputs like accuracy or loss. Log them with mlflow.log_metric(). You can log the same metric many times, which is useful for tracking training curves.

Artifacts. Files such as models, plots, or data snapshots. Log them with mlflow.log_artifact().

Logging Metrics Over Time

Metric history matters for deep learning. You want to see loss drop across epochs. Pass a step argument to plot the curve.

import mlflow
import random

with mlflow.start_run():
    for epoch in range(5):
        # Simulate a decreasing training loss
        loss = 1.0 / (epoch + 1) + random.uniform(-0.05, 0.05)
        mlflow.log_metric("train_loss", loss, step=epoch)
        print(f"Epoch {epoch}: loss={loss:.4f}")

Epoch 0: loss=1.0213
Epoch 1: loss=0.4937
Epoch 2: loss=0.3412
Epoch 3: loss=0.2608
Epoch 4: loss=0.2154

In the UI, open the run and click the "Metrics" tab. MLflow draws the curve for you.

Organizing Runs with Experiments

Create separate experiments for each project. This keeps the dashboard tidy.

import mlflow

# Create or select an experiment by name
mlflow.set_experiment("iris-classification")

with mlflow.start_run(run_name="baseline-lr"):
    mlflow.log_param("model", "logistic_regression")
    mlflow.log_metric("accuracy", 0.97)

You can also add tags to runs. Tags are useful for marking data versions or team names. Use mlflow.set_tag() inside the run.

with mlflow.start_run(run_name="tagged-run"):
    mlflow.set_tag("data_version", "v2.1")
    mlflow.set_tag("owner", "data-team")
    mlflow.log_metric("f1_score", 0.91)

Comparing Runs

The UI makes comparison easy. Select multiple runs and click "Compare". You get a table of parameters and metrics side by side.

You can also do it in code. The search API returns runs as a DataFrame.

import mlflow

# Query all runs in the experiment
runs = mlflow.search_runs(experiment_names=["iris-classification"])

# Show only the columns we care about
print(runs[["tags.mlflow.runName", "metrics.accuracy"]])

  tags.mlflow.runName  metrics.accuracy
0          baseline-lr              0.97

This is powerful for automated model selection. Pick the run with the highest metric and promote it.

Logging Artifacts and Plots

Artifacts are any files you want to keep. Save confusion matrices, feature importance charts, or raw logs.

import matplotlib.pyplot as plt
import mlflow

with mlflow.start_run():
    # Create a simple plot
    plt.plot([1, 2, 3], [3, 2, 1])
    plt.title("Sample Plot")
    plt.savefig("plot.png")

    # Log the file as an artifact
    mlflow.log_artifact("plot.png", artifact_path="plots")

Artifacts appear in the "Artifacts" tab of the run. Large files go to the configured artifact store, such as S3 or a local folder.

Loading a Logged Model

Logged models are reusable. Load them later for inference or deployment.

import mlflow.pyfunc

# Load the model from a specific run URI
model = mlflow.pyfunc.load_model("runs://model")

# Predict on new data
predictions = model.predict([[5.1, 3.5, 1.4, 0.2]])
print(predictions)

[0]

Replace <RUN_ID> with the ID shown in the UI. This makes every experiment a deployable candidate.

Autologging: Less Code, Same Tracking

MLflow can log automatically. Call mlflow.autolog() before training. It captures parameters, metrics, and models for many libraries.

import mlflow
from sklearn.ensemble import RandomForestClassifier
from sklearn.datasets import load_iris

# Enable automatic logging for scikit-learn
mlflow.autolog()

X, y = load_iris(return_X_y=True)

with mlflow.start_run():
    clf = RandomForestClassifier(n_estimators=100)
    clf.fit(X, y)  # Parameters and metrics are logged automatically

Autologging supports scikit-learn, XGBoost, PyTorch, TensorFlow, and more. It is a fast way to start tracking without extra code.

Best Practices for Experiment Tracking

Follow a few simple rules and your tracking stays clean.

Name your runs. Use run_name so results are readable. "lr-0.01-bs-32" beats "run-17".

Log parameters first. Record all inputs before training starts. This keeps the run record consistent.

Keep experiments focused. One experiment per project or model family. Do not mix unrelated work.

Version your data. Add a tag with the dataset version. Link to the data pipeline if needed. If your data flows through message systems, see how to install kafka-python in Python for reliable streaming.

Store artifacts remotely. Use S3 or GCS for team projects. Local folders do not scale.

Common Pitfalls

Beginners often hit the same issues. Here are the fixes.

Runs not showing in the UI. The UI reads from the default store. If you set a custom tracking URI in code, the UI must point to the same place. Set MLFLOW_TRACKING_URI consistently.

Nested runs fail. Use nested=True in mlflow.start_run() when starting a run inside another run.

Missing metrics. Metrics must be numbers. Passing a string raises an error. Convert values with float() first.

Slow logging. Batch metric logging with mlflow.log_metrics() instead of calling the singular version in a tight loop.

Scaling to a Team

A single laptop works for solo projects. Teams need a shared server. Run MLflow with a database backend and shared artifact store.

Use PostgreSQL or MySQL for the backend store. Use S3 or a network drive for artifacts. Point every team member's code to the server URI.

import mlflow

# Point all logging to the shared server
mlflow.set_tracking_uri("http://mlflow.internal:5000")
mlflow.set_experiment("team-recommender")

Now every run lands in one place. Comparison and review become simple.

Connecting to Data Sources

Experiments often read from databases. Log the query version as a tag so results stay traceable. If your features come from an Oracle warehouse, you can query an Oracle database from Python and log the result as an artifact.

For other stacks, the same pattern applies. Log the source, the version, and the extraction time. Future you will thank present you.

Conclusion

MLflow turns messy experimentation into a clean, searchable record. You log parameters, metrics, and artifacts with a few lines of code. The UI lets you compare runs and find the best model fast.

Start small. Install MLflow, run the tracking UI, and wrap your training code in mlflow.start_run(). Add autologging when you want less boilerplate. Move to a shared server when your team grows.

With consistent tracking, every experiment becomes a reusable asset. Your models get better, your results get reproducible, and your team moves faster.