Skip to content

MLflow

MLflow is an open-source platform for managing the ML lifecycle. This sheet covers experiment tracking, model registry, serving, projects, and integrations.

Table of Contents

Installation / Setup

pip install mlflow

# Start MLflow UI
mlflow ui

# Start MLflow server (for teams)
mlflow server --host 0.0.0.0 --port 5000

import mlflow
import mlflow.sklearn
import mlflow.pytorch
import mlflow.pyfunc

# Set tracking URI (local or remote)
mlflow.set_tracking_uri("http://localhost:5000")

# Set experiment
mlflow.set_experiment("my_experiment")

Core Syntax

Experiment Tracking

import mlflow

# Start a run
with mlflow.start_run(run_name="my_run"):
    # Log parameters
    mlflow.log_param("learning_rate", 0.001)
    mlflow.log_param("batch_size", 32)
    mlflow.log_param("model_type", "random_forest")

    # Log metrics
    mlflow.log_metric("accuracy", 0.95)
    mlflow.log_metric("f1_score", 0.93)
    mlflow.log_metric("training_loss", 0.05)

    # Log multiple metrics at once
    mlflow.log_metrics({
        "accuracy": 0.95,
        "precision": 0.92,
        "recall": 0.94,
    })

    # Log artifacts
    mlflow.log_artifact("confusion_matrix.png")
    mlflow.log_artifact("feature_importance.png")

    # Log artifact from directory
    mlflow.log_artifacts("plots/")

Auto-logging

# scikit-learn
mlflow.sklearn.autolog()
model = RandomForestClassifier()
model.fit(X_train, y_train)  # automatically logged

# PyTorch
mlflow.pytorch.autolog()
model.fit(train_loader, epochs=10)  # automatically logged

# Transformers
import transformers
mlflow.transformers.autolog()

# XGBoost
mlflow.xgboost.autolog()

# LightGBM
mlflow.lightgbm.autolog()

Model Logging

# scikit-learn
mlflow.sklearn.log_model(model, "model")

# PyTorch
mlflow.pytorch.log_model(model, "model")

# TensorFlow / Keras
mlflow.tensorflow.log_model(model, "model")

# XGBoost
mlflow.xgboost.log_model(model, "model")

# Generic (pyfunc)
mlflow.pyfunc.log_model(model, "model")

# Log model with signature
from mlflow.types import Schema, ColSpec

input_schema = Schema([
    ColSpec("double", "feature1"),
    ColSpec("double", "feature2"),
])
output_schema = Schema([ColSpec("double", "prediction")])

mlflow.sklearn.log_model(
    model, "model",
    signature=ModelSignature(input_schema, output_schema)
)

Common Operations

Model Registry

# Register model from run
model_uri = "runs:/run_id/model"
mlflow.register_model(model_uri, "my_model")

# Register model from artifact path
mlflow.register_model("runs:/run_id/model", "my_model")

# Get model version
client = mlflow.tracking.MlflowClient()
version = client.get_latest_versions("my_model", stages=["None"])

# Transition model stage
client.transition_model_version_stage(
    name="my_model",
    version=1,
    stage="staging"
)

# Transition to production
client.transition_model_version_stage(
    name="my_model",
    version=1,
    stage="production"
)

# Get model from registry
model = mlflow.pyfunc.load_model("models:/my_model/production")
model = mlflow.sklearn.load_model("models:/my_model/1")

Model Serving

# Serve model as REST API
# mlflow models serve -m "models:/my_model/production" -p 5001

# Predict from served model
import requests
import json

data = {"inputs": [[1.0, 2.0, 3.0]]}
response = requests.post("http://localhost:5001/invocations", json=data)
print(response.json())

Searching Experiments

client = mlflow.tracking.MlflowClient()

# Search runs
runs = client.search_runs(
    experiment_ids=["0"],
    filter_string="metrics.accuracy > 0.9",
    order_by=["metrics.accuracy DESC"],
    max_results=10
)

for run in runs:
    print(f"Run: {run.info.run_id}, Accuracy: {run.data.metrics['accuracy']}")

Comparing Runs

# Get run details
run = client.get_run("run_id")
print(run.data.params)
print(run.data.metrics)

# Compare multiple runs
runs = client.search_runs(experiment_ids=["0"])
for run in runs:
    print(f"Run {run.info.run_id}: {run.data.metrics}")

Practical Examples

Example 1 — Full MLflow workflow

import mlflow
import mlflow.sklearn
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split, GridSearchCV
from sklearn.metrics import accuracy_score

# Set experiment
mlflow.set_experiment("churn_prediction")

# Data
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2)

# Grid search with tracking
param_grid = {
    "n_estimators": [50, 100, 200],
    "max_depth": [5, 10, 20],
}

for n_est in param_grid["n_estimators"]:
    for depth in param_grid["max_depth"]:
        with mlflow.start_run(run_name=f"rf_{n_est}_{depth}"):
            # Train
            model = RandomForestClassifier(n_estimators=n_est, max_depth=depth)
            model.fit(X_train, y_train)

            # Evaluate
            y_pred = model.predict(X_test)
            accuracy = accuracy_score(y_test, y_pred)

            # Log
            mlflow.log_param("n_estimators", n_est)
            mlflow.log_param("max_depth", depth)
            mlflow.log_metric("accuracy", accuracy)
            mlflow.sklearn.log_model(model, "model")

            print(f"n_estimators={n_est}, depth={depth}, accuracy={accuracy:.4f}")

Example 2 — Register best model

import mlflow

client = mlflow.tracking.MlflowClient()

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

# Register best model
model_uri = f"runs:/{best_run.info.run_id}/model"
mlflow.register_model(model_uri, "churn_model")

# Transition to staging
client.transition_model_version_stage("churn_model", 1, "staging")

# Add description
client.update_model_version(
    name="churn_model",
    version=1,
    description=f"Best model from run {best_run.info.run_id}"
)

Example 3 — Load and predict

import mlflow
import pandas as pd

# Load model from registry
model = mlflow.pyfunc.load_model("models:/churn_model/production")

# Predict
data = pd.DataFrame({
    "feature1": [1.0, 2.0, 3.0],
    "feature2": [4.0, 5.0, 6.0],
})
predictions = model.predict(data)
print(predictions)

Common Patterns

Pattern: experiment comparison

client = mlflow.tracking.MlflowClient()
runs = client.search_runs(experiment_ids=["0"])

# Sort by metric
sorted_runs = sorted(runs, key=lambda x: x.data.metrics.get("accuracy", 0), reverse=True)

for run in sorted_runs[:5]:
    print(f"Run {run.info.run_id}: accuracy={run.data.metrics.get('accuracy', 0):.4f}")

Pattern: reproducible runs

import mlflow
import random
import numpy as np
import torch

# Set seeds for reproducibility
random.seed(42)
np.random.seed(42)
torch.manual_seed(42)

# Log seed as parameter
with mlflow.start_run():
    mlflow.log_param("seed", 42)
    # ... training code

Pattern: model signature

from mlflow.models import ModelSignature
from mlflow.types import Schema, ColSpec, DataType

# Define schema
input_schema = Schema([
    ColSpec(DataType.double, "age"),
    ColSpec(DataType.double, "income"),
    ColSpec(DataType.string, "category"),
])
output_schema = Schema([ColSpec(DataType.double, "prediction")])

signature = ModelSignature(inputs=input_schema, outputs=output_schema)

# Log with signature
mlflow.sklearn.log_model(model, "model", signature=signature)

Best Practices

  • Use autologging — saves time, captures all params.
  • Name your runs — makes comparison easier.
  • Log everything — params, metrics, artifacts, data versions.
  • Use model registry — for version control and stage transitions.
  • Set tracking URI — for team collaboration.
  • Use signatures — validates input/output schemas.
  • Tag runs — add context like git commit, dataset version.
  • Compare experiments — use MLflow UI to find best models.

Common Mistakes

Mistake Fix
Not logging parameters Use mlflow.log_param() or autolog
Not tracking experiments Always use mlflow.start_run()
Hardcoding paths Use MLflow's artifact store
Not registering models Use model registry for production
Ignoring reproducibility Log seeds, data versions, code versions

Troubleshooting

Symptom Cause / Fix
Experiments not appearing Check tracking URI; check experiment name
Model not found Check model name; check version/stage
Artifacts not logged Check artifact path; check file permissions
Slow logging Use async logging; batch metric logging
Cannot load model Check model format; check dependencies

Official Documentation

Last Updated

2026-09-06 — Initial version.