ML Experiment Tracking with MLflow

Why Experiment Tracking Matters

Before we jump into code, let's be precise about what experiment tracking is actually solving, because the benefits are deeper than "keeping better notes."

The first benefit is reproducibility. In regulated industries, finance, healthcare, autonomous vehicles, you may be legally required to explain and reproduce model decisions. Even outside those industries, reproducibility is what separates a credible ML claim from a lucky accident. When every run captures the exact parameters, code version, data hash, and environment details, you can recreate any result exactly. Not approximately. Exactly. That's a superpower when auditors, colleagues, or your future self comes asking questions.

The second benefit is collaboration velocity. When multiple engineers are working on the same problem, the default state is duplicated effort. Person A runs depth-5 random forests while Person B runs depth-10 random forests, and they don't discover this until the weekly sync. With a shared tracking server, both can instantly see what's been tried, what the results were, and where the unexplored territory is. Your team stops re-running the same experiments and starts building on each other's work.

The third benefit is debugging efficiency. When your model performance unexpectedly drops, the question is: what changed? If you have a complete record of every run, parameters, metrics, artifacts, data versions, you can bisect the problem. The run before the drop had 91% accuracy with these settings. The run after had 87% with those settings. Something in that delta is your bug. Without tracking, that diagnosis is guesswork. With it, it's a structured investigation.

The fourth benefit is model governance. Organizations that deploy models at scale need to know which model version is live, when it was promoted, who approved it, and what its performance characteristics are. MLflow's Model Registry provides exactly this audit trail. It's the handshake between your research workflow and your production deployment pipeline.

MLflow Architecture

Understanding MLflow's architecture helps you use it correctly and troubleshoot it confidently.

MLflow consists of four core components that work together. The Tracking Server is the heart of the system, it receives logging calls from your training code, stores parameters and metrics in a backend store (SQLite by default, PostgreSQL or MySQL for production), and stores artifacts (model files, plots, datasets) in an artifact store (local filesystem by default, S3 or GCS for production). The MLflow UI is a web application that connects to the Tracking Server and provides the visual interface for browsing experiments, comparing runs, and inspecting artifacts.

The Model Registry is a separate layer on top of the Tracking Server that provides model versioning, lifecycle management, and promotion workflows. It answers the question "which version of this model is in production?" in a structured, auditable way. Finally, MLflow Models is a standard format for packaging trained models with their dependencies so they can be deployed consistently across different serving platforms, REST API, batch inference, Spark, or cloud services like AWS SageMaker.

When you call mlflow.log_metric() in your training script, here's what happens: your code sends an HTTP request to the Tracking Server (or writes directly to a local store if you're not using a remote server). The server writes the metric with a timestamp to the backend store. The artifact store handles any file uploads. The UI reads from these stores when you navigate to a run's page. The whole system is stateless from the client's perspective, your training script just fires off log calls and doesn't need to manage any state itself.

For local development, everything runs on your machine using SQLite and the local filesystem. For team deployments, you stand up a central Tracking Server with a real database and cloud storage, and every team member points their MLFLOW_TRACKING_URI environment variable at it. The code doesn't change; only the destination changes.

MLflow Core Concepts

MLflow organizes experiments in a hierarchy. Let's understand the building blocks before we touch any code.

Experiments and Runs

An experiment is a project (like "fraud detection model"). Within an experiment, you run multiple runs, each one is a single training session with specific parameters and outputs. An experiment is roughly equivalent to "the thing we're trying to figure out," while a run is a single attempt at figuring it out.

Think of it like this:

Experiment: "Customer Churn Prediction"
├── Run 1: Random Forest, max_depth=5, accuracy=0.87
├── Run 2: Random Forest, max_depth=10, accuracy=0.91
├── Run 3: Gradient Boosting, learning_rate=0.01, accuracy=0.93
└── Run 4: Gradient Boosting, learning_rate=0.05, accuracy=0.90

Each run is a snapshot of a training session. Each experiment groups related runs together.

Parameters, Metrics, and Artifacts

Every run can log three types of data. Parameters are the static hyperparameters you pass into training, learning rate, batch size, kernel type, regularization strength. These are the inputs to your training process, and they don't change during a run. Metrics are the performance measurements you get out, accuracy, loss, precision, recall, AUC. These are the outputs, and MLflow lets you log them multiple times per run (for example, after each training epoch) so you can track how they evolve over time. Artifacts are the files your training produces, the trained model itself, plots, preprocessed datasets, evaluation reports. Artifacts are stored as files and can be downloaded or loaded directly through the MLflow API.

Getting Started: Basic Logging

Let's see MLflow in action. First, install it:

pip install mlflow scikit-learn

This single command pulls in MLflow's core libraries, the UI server, and the scikit-learn integration. You don't need to configure anything else for a local setup, it will create an mlruns/ directory in your working folder to store everything.

Here's a before-and-after comparison. First, code without tracking:

from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import accuracy_score
 
# Load data
iris = load_iris()
X_train, X_test, y_train, y_test = train_test_split(
    iris.data, iris.target, test_size=0.2, random_state=42
)
 
# Train
model = RandomForestClassifier(n_estimators=100, max_depth=5)
model.fit(X_train, y_train)
 
# Evaluate
y_pred = model.predict(X_test)
acc = accuracy_score(y_test, y_pred)
print(f"Accuracy: {acc}")

This works, but you've learned nothing that's reproducible. That accuracy number lives only in your terminal history. Now, with MLflow:

import mlflow
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import accuracy_score
 
# Set experiment name (creates it if it doesn't exist)
mlflow.set_experiment("Iris Classification")
 
# Start a run
with mlflow.start_run():
    # Load data
    iris = load_iris()
    X_train, X_test, y_train, y_test = train_test_split(
        iris.data, iris.target, test_size=0.2, random_state=42
    )
 
    # Log parameters
    params = {"n_estimators": 100, "max_depth": 5, "random_state": 42}
    mlflow.log_params(params)
 
    # Train
    model = RandomForestClassifier(**params)
    model.fit(X_train, y_train)
 
    # Evaluate and log metrics
    y_pred = model.predict(X_test)
    acc = accuracy_score(y_test, y_pred)
    mlflow.log_metric("accuracy", acc)
 
    # Save the model
    mlflow.sklearn.log_model(model, "model")
 
    print(f"Accuracy: {acc}")

See? We added maybe 10 lines. Now every parameter, metric, and the model itself are tracked automatically. The with mlflow.start_run(): context manager handles lifecycle, it starts recording everything inside the block and saves it when the block exits. If your script crashes partway through, MLflow marks the run as FAILED rather than leaving ghost data in an ambiguous state.

Digging Deeper: Multi-Metric Tracking and Iteration

Real models need richer logging. When you're doing a grid search or hyperparameter sweep, you want every combination captured as its own run so you can compare them side by side. Here's how to structure that kind of systematic exploration:

import mlflow
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split, cross_val_score
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import accuracy_score, precision_score, recall_score
 
mlflow.set_experiment("Iris Hyperparameter Tuning")
 
# Grid search
depths = [5, 10, 15, 20]
n_estimators = [50, 100, 200]
 
for depth in depths:
    for n_est in n_estimators:
        with mlflow.start_run():
            iris = load_iris()
            X_train, X_test, y_train, y_test = train_test_split(
                iris.data, iris.target, test_size=0.2, random_state=42
            )
 
            # Log hyperparameters
            mlflow.log_param("max_depth", depth)
            mlflow.log_param("n_estimators", n_est)
 
            # Train
            model = RandomForestClassifier(n_estimators=n_est, max_depth=depth)
            model.fit(X_train, y_train)
 
            # Evaluate and log multiple metrics
            y_pred = model.predict(X_test)
            mlflow.log_metric("accuracy", accuracy_score(y_test, y_pred))
            mlflow.log_metric("precision", precision_score(y_test, y_pred, average='weighted'))
            mlflow.log_metric("recall", recall_score(y_test, y_pred, average='weighted'))

Each iteration creates a new run. MLflow captures all combinations automatically, and when you open the UI you can sort all twelve runs by accuracy in one click to find your winner. No spreadsheet required.

Auto-Logging: Let MLflow Do the Work

Logging every metric manually gets tedious, especially when you're working with complex models that have dozens of hyperparameters or training loops that generate hundreds of metrics per epoch. MLflow's auto-logging feature handles it for you. With one line of code, MLflow hooks into your training framework and captures everything it knows how to capture.

For scikit-learn:

import mlflow
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier
 
# Enable auto-logging for sklearn
mlflow.sklearn.autolog()
 
mlflow.set_experiment("Iris Auto-Logging")
 
with mlflow.start_run():
    iris = load_iris()
    X_train, X_test, y_train, y_test = train_test_split(
        iris.data, iris.target, test_size=0.2, random_state=42
    )
 
    model = RandomForestClassifier(n_estimators=100, max_depth=5)
    model.fit(X_train, y_train)
    model.score(X_test, y_test)

That's it. MLflow automatically logs all hyperparameters from the model, standard metrics (accuracy, f1, etc.), the model artifact itself, and even the model signature (input/output schema). It's "smart" enough to know what metrics make sense for classification vs. regression, so a RandomForestClassifier gets accuracy and F1 while a RandomForestRegressor gets R-squared and RMSE.

For PyTorch, there's mlflow.pytorch.autolog(). For TensorFlow: mlflow.tensorflow.autolog(). The pattern is the same across frameworks. Auto-logging is the right default for most projects, you get comprehensive coverage with zero boilerplate, and you can always add additional manual log_metric() calls for custom metrics that the auto-logger doesn't know about.

Logging Artifacts: Models and Beyond

Sometimes you need to save more than just metrics. Confusion matrices, feature importance plots, sample predictions, evaluation reports, preprocessed data files, these are all valuable artifacts that complete the story of a run. Here's how to capture them:

import mlflow
import matplotlib.pyplot as plt
from sklearn.metrics import confusion_matrix, ConfusionMatrixDisplay
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier
 
mlflow.set_experiment("Iris with Artifacts")
 
with mlflow.start_run():
    iris = load_iris()
    X_train, X_test, y_train, y_test = train_test_split(
        iris.data, iris.target, test_size=0.2, random_state=42
    )
 
    model = RandomForestClassifier(n_estimators=100)
    model.fit(X_train, y_train)
 
    # Log the model
    mlflow.sklearn.log_model(model, "model")
 
    # Log metrics
    mlflow.log_metric("accuracy", model.score(X_test, y_test))
 
    # Create and log a confusion matrix plot
    y_pred = model.predict(X_test)
    cm = confusion_matrix(y_test, y_pred)
 
    fig, ax = plt.subplots()
    disp = ConfusionMatrixDisplay(confusion_matrix=cm, display_labels=iris.target_names)
    disp.plot(ax=ax)
    plt.savefig("confusion_matrix.png")
 
    mlflow.log_artifact("confusion_matrix.png")
 
    # Or log raw data as JSON
    mlflow.log_dict({"class_names": list(iris.target_names)}, "metadata.json")

Now your run has the trained model (loadable via mlflow.sklearn.load_model()), the confusion matrix PNG (visible in the MLflow UI), metadata JSON (for reference), and all metrics and parameters. A colleague can open this run six months later, see every detail at a glance, and load the model directly from the artifact store, no file hunting, no guessing.

The MLflow UI: Seeing Your Experiments

Logging is great, but you need to see what you've logged. MLflow comes with a web UI that turns raw run data into an interactive comparison interface.

In your project directory, run:

mlflow ui

This starts a local server (usually at http://localhost:5000). Open it in your browser and you'll see an Experiments tab with all your experiments listed, a Runs tab showing all runs within an experiment with their metrics and parameters, interactive Metrics charts that let you visualize how metrics changed across runs or over training time, a sortable Parameters table for comparing parameter values side-by-side, and an Artifacts panel where you can browse and download logged files including model files and plots.

The UI is where experiment tracking shines. You can instantly sort all runs by accuracy to find your best performer, filter by parameter values to understand sensitivity, and visualize training curves from multiple runs on the same chart. What used to take twenty minutes of script-wrangling and spreadsheet work becomes a two-second click.

Model Registry Workflow

The Model Registry is where MLflow goes from "nice experiment logging tool" to "real production ML infrastructure." It's the bridge between your training environment and your deployment environment, and understanding its workflow is essential for teams that want to deploy models responsibly.

Here's how the promotion lifecycle works. A model starts as a logged artifact inside a run, just a file. You then register it with a name, which creates a versioned entry in the Model Registry. From there, the model moves through lifecycle stages: None (just registered), Staging (deployed to a test or pre-production environment for validation), and Production (live). There's also Archived for models that have been retired. Each transition is an explicit, logged event, so you have a complete audit trail of when a model was promoted, who promoted it, and why.

import mlflow
 
# After training your best model
mlflow.set_experiment("Production Models")
 
with mlflow.start_run():
    model = RandomForestClassifier(n_estimators=100, max_depth=10)
    model.fit(X_train, y_train)
 
    # Log the model
    mlflow.sklearn.log_model(model, "model")
 
    # Register it
    model_uri = f"runs:/{mlflow.active_run().info.run_id}/model"
    mv = mlflow.register_model(model_uri, "ChurnPredictionModel")
 
    # Transition to "Staging"
    mlflow.tracking.MlflowClient().transition_model_version_stage(
        name="ChurnPredictionModel",
        version=mv.version,
        stage="Staging"
    )

Now you have a named model with version control, stage management, and full metadata. Your deployment system can always query for the current Production version by name rather than hunting for a specific file path. When you train an improved version, you register it, validate it in Staging, and promote it to Production, all tracked, all auditable, all reversible.

The Model Registry also supports annotations: you can attach descriptions to each version explaining what changed, link it back to the experiment run it came from, and add custom tags for things like "approved-by" or "data-version." This is the infrastructure that makes ML deployment look like software engineering rather than folklore.

Running MLflow Server Remotely

So far, we've run everything locally. But in a team, you want a centralized tracking server that everyone's experiments flow into.

On your server machine, run:

mlflow server --host 0.0.0.0 --port 5000 --backend-store-uri postgresql://user:pass@dbhost/mlflow --default-artifact-root s3://my-bucket/mlflow-artifacts

This uses a PostgreSQL database (for scalability and concurrent access) and S3 (for artifact storage). Now all your team's experiments are centralized, and the UI shows everyone's runs in one place.

From your local machine, point your code at the server:

import mlflow
 
mlflow.set_tracking_uri("http://your-server:5000")
mlflow.set_experiment("Team Experiment")
 
with mlflow.start_run():
    # Your code...
    model.fit(X_train, y_train)
    mlflow.log_metric("accuracy", model.score(X_test, y_test))

Every run uploads to the central server. No confusion about whose experiment is whose. You can also set the tracking URI via the MLFLOW_TRACKING_URI environment variable instead of hardcoding it, which makes your code environment-agnostic, the same script works locally and in CI without modification.

Integrating MLflow Into Existing Code

You're thinking: "This sounds great, but our codebase is huge. Will integration be painful?" The honest answer is: almost never. MLflow is designed to be minimally invasive. You wrap your existing training code in a start_run() context and add log calls, you don't restructure anything. Here's a production pipeline before:

def train_pipeline(config_file):
    config = load_config(config_file)
    X_train, y_train = load_data(config['train_path'])
    X_test, y_test = load_data(config['test_path'])
 
    model = train_model(X_train, y_train, config)
    metrics = evaluate_model(model, X_test, y_test)
 
    save_model(model, 'outputs/model.pkl')
    print(f"Metrics: {metrics}")
 
train_pipeline('config.yaml')

And after:

def train_pipeline(config_file):
    import mlflow
 
    mlflow.set_experiment("Production Pipeline")
 
    with mlflow.start_run():
        config = load_config(config_file)
        mlflow.log_params(config)  # Log the entire config
 
        X_train, y_train = load_data(config['train_path'])
        X_test, y_test = load_data(config['test_path'])
 
        model = train_model(X_train, y_train, config)
        metrics = evaluate_model(model, X_test, y_test)
 
        mlflow.log_metrics(metrics)
        mlflow.sklearn.log_model(model, 'model')
 
        # Optionally save locally too
        save_model(model, 'outputs/model.pkl')
 
        print(f"Metrics: {metrics}")
 
train_pipeline('config.yaml')

Four lines added. No refactoring. No breaking changes. The function signature is the same, the data flow is the same, and all your downstream code still works. This is by design, the MLflow team knew they were asking people to instrument existing code, so they made it as frictionless as possible.

Common MLflow Mistakes

Even experienced ML practitioners make these mistakes when they start using MLflow. Knowing them in advance will save you the debugging time.

The most common mistake is not naming your experiments. If you never call mlflow.set_experiment(), everything goes into the "Default" experiment. You end up with hundreds of runs from completely unrelated projects all mixed together, and the UI becomes a mess that's harder to navigate than no tracking at all. Always set experiment names that describe what question you're investigating.

The second mistake is logging too coarsely or too finely. If you only log final accuracy, you lose the ability to diagnose training instability. If you log every single batch loss value in a model that trains for 100 epochs with 1000 batches each, you're generating 100,000 metric points per run and the UI becomes slow to load. A good rule of thumb: log per-epoch metrics for training curves, and log final metrics at the end. Use step parameter in log_metric() to track the epoch number.

The third mistake is forgetting to log the data version. You track your code version via git commits and your model hyperparameters via log_params(), but if you forget to log which version of the training data was used, you'll encounter runs that can't be reproduced even when everything else is identical. Log a hash of your training data or the data pipeline version as a tag or parameter.

The fourth mistake is using local artifact stores in team settings. If you stand up a central Tracking Server but forget to configure --default-artifact-root to point to shared storage (S3, GCS, Azure Blob), artifacts get stored on the server's local disk. Every team member can see the run metadata in the UI but can't actually access the artifacts. Always configure shared artifact storage for team setups.

The fifth mistake is ignoring run names. By default, MLflow generates random adjective-noun run names. After fifty runs, you have "laughing-crow-47" and "quirky-emu-12" and no idea what either of them represents. Call mlflow.start_run(run_name="depth10_features_v2") to give meaningful names, or use mlflow.set_tag("description", "Testing new feature engineering pipeline") to add context you can search for later.

Best Practices: Tracking Like a Pro

A few guidelines to keep your experiment tracking clean and actually useful.

Tag your runs for easy filtering:

mlflow.set_tag("team", "data-science")
mlflow.set_tag("feature_set", "v2_with_embeddings")
mlflow.set_tag("status", "baseline")

Log your code version (git commit hash):

import subprocess
commit = subprocess.check_output(['git', 'rev-parse', 'HEAD']).decode().strip()
mlflow.set_tag("git_commit", commit)

Use descriptive experiment names. Not "experiment_1", but "RandomForest_FeatureV2_CV5".

Log early and often. Don't wait until the end of a long training. Log after each epoch or batch if it helps you debug.

Organize artifacts in subdirectories:

mlflow.log_artifact("plots/confusion_matrix.png", "visualizations")
mlflow.log_artifact("data/class_distribution.json", "data-summary")

Summary

Experiment tracking isn't a luxury, it's the difference between doing ML and doing ML professionally. The chaos of untracked experiments isn't just annoying; it's expensive. It costs you time when you can't reproduce a result. It costs your team velocity when everyone re-runs experiments that have already been tried. It costs you credibility when you can't explain how a deployed model was built.

MLflow removes all of those costs with minimal overhead. You get automatic logging of every parameter, metric, and artifact. You get a powerful visual UI that turns raw run data into instant insights. You get a Model Registry that turns your best experiment results into versioned, promotable, auditable production assets. And you get integrations with every major ML framework so you can capture all of this with a single autolog() call if you want to go fast.

Start small. Add mlflow.log_metric() and mlflow.set_experiment() to one script. Run mlflow ui and watch your experiments appear in the browser. Once you see your first side-by-side comparison of two runs, instantly understanding which hyperparameter combination won and by exactly how much, you'll feel the cognitive load lift. You'll stop keeping mental notes and start trusting the system. That's when the real productivity gains kick in.

Your future self will thank you when you need to reproduce a result six months from now. Your team will thank you when they can build on your experiments instead of re-running them. And your stakeholders will thank you when you can explain exactly how your production model was trained, validated, and promoted, with a complete audit trail to back it up.

This is what professional ML looks like. Now go track something.