MLflow: Experiment Tracking, Model Registry & Serving

What is MLflow?

MLflow is an open-source platform created by Databricks (2018) that solves the reproducibility and management problem in ML: "which parameters produced my best model, and how do I get it back to production?" It provides a clean API for logging any training run, storing artifacts, and managing model versions through a standardized lifecycle.

MLflow is framework-agnostic — the same tracking API works whether you are training a scikit-learn pipeline, a PyTorch transformer, or a custom NumPy model. Runs are stored in a backend (local filesystem, PostgreSQL, MySQL) and artifacts in an artifact store (local disk, S3, GCS, Azure Blob).

Quickstart: Log Your First Run

Install MLflow with pip install mlflow, then wrap your training code in a run context:

import mlflow
import mlflow.sklearn
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import accuracy_score

mlflow.set_experiment("fraud-detection")

with mlflow.start_run(run_name="rf-baseline"):
    # Log hyperparameters
    params = {"n_estimators": 100, "max_depth": 8, "random_state": 42}
    mlflow.log_params(params)

    # Train
    clf = RandomForestClassifier(**params)
    clf.fit(X_train, y_train)

    # Log metrics
    acc = accuracy_score(y_test, clf.predict(X_test))
    mlflow.log_metric("accuracy", acc)
    mlflow.log_metric("f1_score", f1)

    # Log the model artifact
    mlflow.sklearn.log_model(clf, artifact_path="model")
    print(f"Run ID: {mlflow.active_run().info.run_id}")

Run mlflow ui to open the tracking dashboard at http://localhost:5000. Every run appears as a row with its params, metrics, and artifact links.

Autolog: Zero-Code Tracking

For supported frameworks, one line before training is enough:

import mlflow

# Enable autolog for all supported libraries at once
mlflow.autolog()

# Or per-library for fine-grained control
mlflow.sklearn.autolog(log_model_signatures=True)
mlflow.pytorch.autolog(log_every_n_epoch=5)
mlflow.xgboost.autolog()
mlflow.transformers.autolog()

# Training code unchanged — MLflow captures everything automatically
from sklearn.linear_model import LogisticRegression
model = LogisticRegression(C=0.5, max_iter=200)
model.fit(X_train, y_train)
# ✓ Params logged: C=0.5, max_iter=200, solver=lbfgs
# ✓ Metrics logged: training_score, validation metrics
# ✓ Model artifact saved automatically

Model Registry: Staging → Production Lifecycle

Register a model from a run, then manage its lifecycle programmatically:

from mlflow.tracking import MlflowClient

client = MlflowClient()

# Register model from a completed run
run_id = "abc123def456"
model_uri = f"runs:/{run_id}/model"
mv = client.create_model_version(
    name="FraudDetector",
    source=model_uri,
    run_id=run_id
)
print(f"Version: {mv.version}")  # Version: 3

# Transition through lifecycle stages
client.transition_model_version_stage(
    name="FraudDetector",
    version=3,
    stage="Staging",
    archive_existing_versions=False,
)

# After validation, promote to Production
client.transition_model_version_stage(
    name="FraudDetector",
    version=3,
    stage="Production",
    archive_existing_versions=True,  # archive old Production version
)

# Load production model anywhere
prod_model = mlflow.pyfunc.load_model("models:/FraudDetector/Production")

MLflow vs Weights & Biases vs Comet ML