Manage Your Machine Learning Lifecycle with MLflow in Python

If you have trained more than three machine learning models, you already know the chaos: folders named final_model, final_model_v2, final_model_THIS_ONE, parameter files scattered across notebook comments, and the nagging question of which hyperparameters you used in that model that worked so well two weeks ago. MLflow exists precisely to solve that problem. In this article you'll track experiments with parameters, metrics, and artifacts, query the best run programmatically, and promote a model to production with the Model Registry — all in a few lines of Python.

The Chaos Without Experiment Tracking

As an ML project grows, the lack of a tracking system creates concrete problems:

• Broken reproducibility: you cannot recreate a result because you don't remember which data version you used or what value max_depth had.
• Impossible comparison: comparing 15 runs manually by reviewing scattered notes is error-prone.
• Difficult collaboration: a teammate cannot pick up your work if artifacts and parameters only exist on your machine.
• Lost knowledge: the model you discarded as "bad" might have been the best for a different use case, but you have no way to recover it.

A tracking system is not a luxury; it is basic engineering hygiene for serious ML projects.

MLflow: An Overview

MLflow is an open-source platform created by Databricks to manage the complete ML lifecycle. It is divided into four main components:

• MLflow Tracking: records and queries experiments — parameters, metrics, artifacts, and models.
• MLflow Projects: packages ML code in a reproducible format that can run on any platform.
• MLflow Models: defines a standard format for packaging models with support for multiple "flavors" (sklearn, tensorflow, pytorch, etc.).
• MLflow Model Registry: a centralized model store with versioning, staging, and annotations.

For most individual or small-team projects, you will start with Tracking and, when you have models in production, add the Registry.

Installation and Setup

Installation is straightforward:

pip install mlflow scikit-learn pandas numpy

To launch the local web interface:

mlflow ui

This starts a server at http://127.0.0.1:5000 where you can visualize all your experiments. By default, MLflow stores data in an mlruns/ directory in your current working directory.

Tracking Experiments with Python

We will use the Iris dataset and a Random Forest to illustrate all the tracking features. The goal is to train several configurations and automatically log everything.

import mlflow
import mlflow.sklearn
import numpy as np
import pandas as pd
from sklearn.datasets import load_iris
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score, f1_score, precision_score, recall_score

# Load data
iris = load_iris()
X = pd.DataFrame(iris.data, columns=iris.feature_names)
y = iris.target

X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42
)

# Set experiment name
mlflow.set_experiment("iris-random-forest")

# Configurations to try
configs = [
    {"n_estimators": 50,  "max_depth": 3, "min_samples_split": 2},
    {"n_estimators": 100, "max_depth": 5, "min_samples_split": 5},
    {"n_estimators": 200, "max_depth": 10, "min_samples_split": 2},
]

for config in configs:
    with mlflow.start_run():
        # Log parameters
        mlflow.log_param("n_estimators", config["n_estimators"])
        mlflow.log_param("max_depth", config["max_depth"])
        mlflow.log_param("min_samples_split", config["min_samples_split"])
        mlflow.log_param("test_size", 0.2)
        mlflow.log_param("random_state", 42)

        # Train model
        model = RandomForestClassifier(
            n_estimators=config["n_estimators"],
            max_depth=config["max_depth"],
            min_samples_split=config["min_samples_split"],
            random_state=42
        )
        model.fit(X_train, y_train)

        # Evaluate
        y_pred = model.predict(X_test)
        acc  = accuracy_score(y_test, y_pred)
        f1   = f1_score(y_test, y_pred, average="weighted")
        prec = precision_score(y_test, y_pred, average="weighted")
        rec  = recall_score(y_test, y_pred, average="weighted")

        # Log metrics
        mlflow.log_metric("accuracy", acc)
        mlflow.log_metric("f1_score", f1)
        mlflow.log_metric("precision", prec)
        mlflow.log_metric("recall", rec)

        # Log feature importance as artifact
        feature_importance = pd.DataFrame({
            "feature": iris.feature_names,
            "importance": model.feature_importances_
        }).sort_values("importance", ascending=False)
        feature_importance.to_csv("feature_importance.csv", index=False)
        mlflow.log_artifact("feature_importance.csv")

        # Save the model with an input example for automatic documentation
        mlflow.sklearn.log_model(
            model,
            "random_forest_model",
            input_example=X_train.iloc[:3]
        )

        print(f"Run complete — accuracy: {acc:.4f}, f1: {f1:.4f}")

Here is what each key call does:

• mlflow.set_experiment("iris-random-forest"): groups all runs under the same experiment. If it does not exist, MLflow creates it.
• mlflow.start_run(): opens a run context. When the with block exits, the run closes automatically.
• mlflow.log_param(key, value): records a hyperparameter or configuration (typically values that do not change during training).
• mlflow.log_metric(key, value): records an evaluation metric (can be called multiple times with step to plot training curves).
• mlflow.log_artifact(path): uploads any file to the artifact store.
• mlflow.sklearn.log_model(model, name): serializes and saves the model with enough metadata to load it with any compatible flavor.

The MLflow UI

After running the script, open http://127.0.0.1:5000 in your browser. You will see:

• Experiment list in the left panel.
• Runs table with all runs, showing parameters and metrics in configurable columns.
• Filters and sorting to quickly find the best run by any metric.
• Individual run view with all parameters, metrics, downloadable artifacts, and the model URI.

The UI also renders metric charts over time when you log them with step, which is very useful for iterative training (epochs in neural networks, for example).

Comparing Runs and Finding the Best Model

You can compare runs from the UI by selecting several and clicking "Compare". You can also do it programmatically:

from mlflow.tracking import MlflowClient

client = MlflowClient()

# Get the experiment by name
experiment = client.get_experiment_by_name("iris-random-forest")
experiment_id = experiment.experiment_id

# Find the best run by accuracy
best_run = client.search_runs(
    experiment_ids=[experiment_id],
    order_by=["metrics.accuracy DESC"],
    max_results=1
)[0]

print(f"Best run ID: {best_run.info.run_id}")
print(f"Accuracy: {best_run.data.metrics['accuracy']:.4f}")
print(f"Parameters: {best_run.data.params}")

Loading a Logged Model and Making Predictions

Once the best run is identified, loading the model is straightforward:

import mlflow.sklearn
import pandas as pd

# Use the run_id from the best run
run_id = best_run.info.run_id
model_uri = f"runs:/{run_id}/random_forest_model"

# Load the model
loaded_model = mlflow.sklearn.load_model(model_uri)

# New data to predict
new_data = pd.DataFrame({
    "sepal length (cm)": [5.1, 6.7],
    "sepal width (cm)":  [3.5, 3.0],
    "petal length (cm)": [1.4, 5.2],
    "petal width (cm)":  [0.2, 2.3]
})

predictions = loaded_model.predict(new_data)
species = ["setosa", "versicolor", "virginica"]
print([species[p] for p in predictions])
# ['setosa', 'virginica']

Registering the Model in the Model Registry

For teams managing production models, the Registry adds a governance layer. MLflow 3 replaced the stage system (Production, Staging, Archived) with model aliases, which are more flexible: you can have multiple aliases per version and reassign them without state transition restrictions.

# Register the best model
result = mlflow.register_model(
    model_uri=f"runs:/{run_id}/random_forest_model",
    name="IrisClassifier"
)

# Assign the "champion" alias to the production version
client.set_registered_model_alias(
    name="IrisClassifier",
    alias="champion",
    version=result.version
)

# Always load the version with the champion alias
prod_model = mlflow.sklearn.load_model("models:/IrisClassifier@champion")

With this setup, your inference pipeline always points to the @champion alias without hardcoding run IDs or version numbers. When you retrain and have a better model, you simply move the alias to the new version.

Conclusion

MLflow transforms the way you work with machine learning: from a collection of disorganized scripts and files to a reproducible, auditable workflow. With just four calls (log_param, log_metric, log_artifact, log_model) you get complete experiment tracking.

The natural next step is to integrate MLflow into a CI/CD pipeline: run automatic training when new data arrives, log the run, and automatically promote the model to the Registry if it surpasses a quality threshold. Tools like GitHub Actions or Azure Pipelines integrate well with the MLflow API for this purpose.

The code in this article is designed to run directly; you only need to install the dependencies and you will have a functional experimentation system in minutes.

Once your experiments are tracked, the next step is wrapping your training code in a reliable scheduled workflow: Orchestrate ML Pipelines with Prefect. And when your @champion model is ready to receive traffic, see Serve a Machine Learning Model with FastAPI and Docker.