1 The Problem: Tracking Chaos at Scale

By Lesson 63 you've run GridSearchCV sweeps, RandomizedSearchCV sweeps, and Optuna studies with hundreds of trials. Each of those tools gives you the best hyperparameters at the end of a run — but the moment you close that notebook, restart your kernel, or come back the next day, you've usually lost everything except whatever number you wrote down by hand. Scale that to a real project: ten team members, six months, hundreds of model variants, multiple datasets, several feature engineering pipelines. The result is a familiar mess:

• Spreadsheet tracking: someone manually copies accuracy numbers into a Google Sheet, which inevitably falls out of sync with what's actually in the model file.
• Filename versioning: model.pkl, model_v2.pkl, model_v2_tuned.pkl, model_v2_tuned_FINAL.pkl — with no record of what changed between them or which one is actually deployed.
• Lost reproducibility: six weeks later, nobody can answer "what hyperparameters, what training data snapshot, and what code commit produced the model currently in production?"
• No audit trail: a regulator, auditor, or incident responder asks "who approved this model for production, and based on what evaluation?" — and there's no answer.
• Duplicated work: a colleague reruns an experiment you already tried and discarded three months ago, because there was no shared, searchable record of what had been tried.

MLflow (open-sourced by Databricks in 2018) is the most widely adopted answer to this problem. It is not a single tool but a small platform built from four loosely-coupled components that you can adopt incrementally:

This lesson covers all four, plus the operational question every team eventually hits: should you run MLflow's default local file store, or a centralized server that the whole team points at?

2 MLflow Tracking: Logging Runs

MLflow organizes work into experiments (a named collection of related runs, e.g. "churn-model-xgboost") and runs (one execution — one training job, one hyperparameter configuration). Each run can log:

• Parameters — inputs that don't change during the run: n_estimators=300, learning_rate=0.05
• Metrics — numeric outputs, optionally logged at multiple steps/epochs: val_auc=0.891, train_loss per epoch
• Artifacts — arbitrary files: the trained model, a confusion matrix plot, a SHAP summary plot (Lesson 64), the feature list
• Tags & source metadata — git commit hash, source script name, who ran it

import mlflow
import mlflow.sklearn
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import roc_auc_score, f1_score, accuracy_score
from sklearn.model_selection import train_test_split
import pandas as pd

# Point at a tracking store. Local default: a './mlruns' folder.
# (Section 8 covers pointing this at a centralized server instead.)
mlflow.set_tracking_uri("file:./mlruns")
mlflow.set_experiment("churn-prediction")

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

params = {
    "n_estimators": 300,
    "max_depth": 8,
    "min_samples_leaf": 5,
    "random_state": 42,
}

with mlflow.start_run(run_name="rf-baseline"):
    # 1. Log parameters
    mlflow.log_params(params)
    mlflow.log_param("feature_set", "v3_with_tenure_buckets")

    # 2. Train (exactly as in Lessons 11-20 — nothing changes here)
    model = RandomForestClassifier(**params)
    model.fit(X_train, y_train)

    # 3. Evaluate and log metrics
    val_proba = model.predict_proba(X_test)[:, 1]
    val_pred  = model.predict(X_test)
    mlflow.log_metric("val_auc", roc_auc_score(y_test, val_proba))
    mlflow.log_metric("val_f1", f1_score(y_test, val_pred))
    mlflow.log_metric("val_accuracy", accuracy_score(y_test, val_pred))

    # 4. Log artifacts — any file at all
    feature_importance = pd.Series(
        model.feature_importances_, index=X_train.columns
    ).sort_values(ascending=False)
    feature_importance.to_csv("feature_importance.csv")
    mlflow.log_artifact("feature_importance.csv")

    # 5. Log the model itself in MLflow's standard format (Section 4)
    mlflow.sklearn.log_model(model, artifact_path="model")

    print(f"Run ID: {mlflow.active_run().info.run_id}")

Autologging: Zero-Effort Tracking

Manually calling log_param for every hyperparameter gets tedious and you will inevitably forget one. MLflow's autologging instruments the most common ML libraries so that calling .fit() automatically logs parameters, metrics, and the model artifact — no manual calls needed.

import mlflow
import mlflow.sklearn
from sklearn.ensemble import GradientBoostingClassifier

mlflow.set_experiment("churn-prediction")

# One line. Captures all constructor params, training metrics,
# the fitted model, and a model signature (input/output schema).
mlflow.sklearn.autolog()

with mlflow.start_run(run_name="gbm-autolog"):
    model = GradientBoostingClassifier(
        n_estimators=250, learning_rate=0.03, max_depth=4, random_state=42
    )
    model.fit(X_train, y_train)
    # No explicit log_metric calls needed — autolog captures
    # training score, and (if you call model.score / cross_val_score)
    # those too. You can still log additional custom metrics:
    test_auc = roc_auc_score(y_test, model.predict_proba(X_test)[:, 1])
    mlflow.log_metric("test_auc", test_auc)

# Autologging also exists for deep learning frameworks
import mlflow.pytorch
import torch
import torch.nn as nn

mlflow.set_experiment("churn-prediction-dl")
mlflow.pytorch.autolog(log_models=True, log_every_n_epoch=1)

class ChurnNet(nn.Module):
    def __init__(self, n_features):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(n_features, 64), nn.ReLU(), nn.Dropout(0.3),
            nn.Linear(64, 16), nn.ReLU(),
            nn.Linear(16, 1),
        )
    def forward(self, x):
        return self.net(x)

with mlflow.start_run(run_name="churn-mlp"):
    model = ChurnNet(n_features=X_train.shape[1])
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
    loss_fn = nn.BCEWithLogitsLoss()

    for epoch in range(20):
        optimizer.zero_grad()
        logits = model(torch.tensor(X_train.values, dtype=torch.float32))
        loss = loss_fn(logits.squeeze(), torch.tensor(y_train.values, dtype=torch.float32))
        loss.backward()
        optimizer.step()
        # mlflow.pytorch.autolog logs the optimizer/loss-fn config automatically;
        # per-epoch metrics still need an explicit call inside a training loop
        mlflow.log_metric("train_loss", loss.item(), step=epoch)

3 Comparing Runs in the MLflow UI

Every run you log accumulates in the tracking store. Launch the UI from the same directory as your mlruns folder:

mlflow ui --port 5000
# Then open http://localhost:5000

The UI lets you select multiple runs and compare them side-by-side — sortable by any logged metric, with parallel-coordinates plots to see which hyperparameters correlate with better performance. This is the payoff for the logging discipline in Section 2: instead of scrolling through notebook output, you get a queryable table.

You can also query runs programmatically — useful for automated "pick the best model" steps in a pipeline:

from mlflow.tracking import MlflowClient

client = MlflowClient(tracking_uri="file:./mlruns")
experiment = client.get_experiment_by_name("churn-prediction")

# Search runs, ordered by validation AUC, descending
runs = client.search_runs(
    experiment_ids=[experiment.experiment_id],
    filter_string="metrics.val_auc > 0.88",
    order_by=["metrics.val_auc DESC"],
    max_results=5,
)

print(f"{'Run Name':20s} {'val_auc':>10s} {'val_f1':>10s}")
for run in runs:
    name = run.data.tags.get("mlflow.runName", run.info.run_id[:8])
    print(f"{name:20s} {run.data.metrics.get('val_auc', float('nan')):10.3f} "
          f"{run.data.metrics.get('val_f1', float('nan')):10.3f}")

best_run = runs[0]
print(f"\nBest run: {best_run.info.run_id} (val_auc={best_run.data.metrics['val_auc']:.3f})")

4 MLflow Models: A Standard Packaging Format

A trained scikit-learn model, a PyTorch state_dict, and an XGBoost booster all have different save formats and loading APIs. MLflow Models solves this with a uniform packaging convention: every model logged via mlflow.<library>.log_model() is saved as a directory containing the serialized model plus an MLmodel metadata file describing its flavors — the different ways the model can be loaded and served.

import mlflow.sklearn

with mlflow.start_run(run_name="rf-for-deployment") as run:
    model = RandomForestClassifier(n_estimators=300, max_depth=8, random_state=42)
    model.fit(X_train, y_train)

    # Log with an explicit input/output signature — strongly recommended.
    # The signature documents (and validates) the expected schema at serve time.
    from mlflow.models import infer_signature
    signature = infer_signature(X_train, model.predict(X_train))

    mlflow.sklearn.log_model(
        sk_model=model,
        artifact_path="model",
        signature=signature,
        input_example=X_train.iloc[:5],
    )
    model_uri = f"runs:/{run.info.run_id}/model"
    print(f"Model logged at: {model_uri}")

The "flavor" system is what makes the format universal: the sklearn flavor lets you load it back as a normal RandomForestClassifier, while the python_function ("pyfunc") flavor is a generic, library-agnostic interface that every MLflow model supports — meaning any serving tool that speaks "pyfunc" can serve a model regardless of whether it was trained with scikit-learn, PyTorch, XGBoost, or LightGBM.

import mlflow.pyfunc
import mlflow.sklearn

# Load back using the library-specific flavor (returns the raw estimator)
sk_model = mlflow.sklearn.load_model(model_uri)
print(type(sk_model))  # <class 'sklearn.ensemble._forest.RandomForestClassifier'>

# Or load using the generic pyfunc flavor (works identically for any model type)
pyfunc_model = mlflow.pyfunc.load_model(model_uri)
predictions = pyfunc_model.predict(X_test.iloc[:5])
print(predictions)

Because the packaged model already declares its environment, you can serve it as a local REST API with one command — no need to hand-write a Flask app just to test it:

# Spin up a local REST endpoint for the model — useful for smoke-testing
# before handing it off to the FastAPI/Docker deployment pipeline (Lesson 66)
mlflow models serve -m runs:/7e3f1a2b9c4d4e8f9a0b1c2d3e4f5a6b/model -p 1234 --env-manager local

# Query it like any REST API
curl -X POST http://127.0.0.1:1234/invocations \
  -H "Content-Type: application/json" \
  -d '{"dataframe_split": {"columns": ["tenure_months", "monthly_charges"], "data": [[14, 79.99]]}}'

5 MLflow Projects: Packaging Reproducible Runs

Logging a run tells you what happened. An MLflow Project tells you how to make it happen again — on a teammate's laptop, on a CI runner, or on a cloud training cluster — without them needing to reverse-engineer your environment from a README. A project is just a directory with an MLproject file declaring entry points, their parameters, and the environment to run them in.

# MLproject (no file extension — the filename itself is the convention)
name: churn-prediction

# Reproducible environment — conda.yaml or a Dockerfile both work
conda_env: conda.yaml

entry_points:
  main:
    parameters:
      n_estimators: {type: int, default: 300}
      max_depth: {type: int, default: 8}
      data_path: {type: str, default: "data/churn.csv"}
    command: "python train.py --n_estimators {n_estimators} --max_depth {max_depth} --data_path {data_path}"

  evaluate:
    parameters:
      model_uri: {type: str}
      test_data: {type: str, default: "data/churn_test.csv"}
    command: "python evaluate.py --model_uri {model_uri} --test_data {test_data}"

# conda.yaml — pinned environment for the "main" entry point
name: churn-prediction-env
channels:
  - conda-forge
dependencies:
  - python=3.10
  - pip
  - pip:
      - mlflow==2.13.0
      - scikit-learn==1.4.0
      - pandas==2.2.0

Anyone with MLflow installed can now run the exact same experiment — MLflow handles creating the conda environment and invoking the entry point command:

# Run locally from the project directory
mlflow run . -P n_estimators=500 -P max_depth=10

# Or run directly from a git repo, no manual cloning required —
# this is the "send a teammate one command" reproducibility story
mlflow run https://github.com/your-org/churn-prediction.git \
  -P n_estimators=500 --experiment-name churn-prediction

6 The Model Registry: Versioning and Promotion Workflows

Tracking answers "what did I try?" The Model Registry answers the question that matters once a model is good enough to matter to the business: "which exact version is live right now, and what's its history?" A registered model is a named, versioned entity, independent of any single run, that tracks a lineage of versions as they move through stages.

Registering a Model

import mlflow
from mlflow.tracking import MlflowClient

run_id = "7e3f1a2b9c4d4e8f9a0b1c2d3e4f5a6b"
model_uri = f"runs:/{run_id}/model"

# Registers a new version under the name "churn-classifier"
# (creates the registered model the first time it's called)
result = mlflow.register_model(model_uri=model_uri, name="churn-classifier")
print(f"Registered '{result.name}' as version {result.version}")

Stages: Staging → Production → Archived

Each version of a registered model carries a stage: None (just registered), Staging (under validation), Production (serving live traffic), or Archived (retired). Moving a version between stages is an explicit, logged action — this is the approval gate that gives you an audit trail.

client = MlflowClient()

# Promote version 4 to Staging for QA validation
client.transition_model_version_stage(
    name="churn-classifier",
    version=4,
    stage="Staging",
    archive_existing_versions=False,
)

# ... QA team runs validation checks against the Staging model ...
# ... a human (or an automated gate) reviews metrics and signs off ...

# Promote to Production. archive_existing_versions=True automatically
# demotes whatever was previously in Production to Archived — there is
# always exactly one unambiguous "current production version."
client.transition_model_version_stage(
    name="churn-classifier",
    version=4,
    stage="Production",
    archive_existing_versions=True,
)

# Attach a human-readable note documenting the approval decision
client.update_model_version(
    name="churn-classifier",
    version=4,
    description=(
        "Approved for production 2026-07-01 by M. Alvarez (Lead DS). "
        "val_auc=0.913, beats prior production version (0.887) by 2.6pp. "
        "Fairness check (Lesson 64 SHAP audit) showed no proxy discrimination "
        "via tenure_months."
    ),
)

Loading the Production Model by Alias

Downstream consumers — a batch scoring job, the FastAPI service from Lesson 66 — never hardcode a run ID. They reference the registry by name and stage, so promoting a new version requires zero code changes downstream:

import mlflow.pyfunc

# Always resolves to whichever version currently holds the Production stage
production_model = mlflow.pyfunc.load_model("models:/churn-classifier/Production")
predictions = production_model.predict(new_customers_df)

7 Centralized Tracking Server vs. Local File Store

Everything so far has used mlflow.set_tracking_uri("file:./mlruns") — the default, which writes run metadata as files and folders directly on your machine. That's fine for solo experimentation, but it breaks down the moment a team needs to share results: there's no central place to compare a colleague's runs against yours, and the Model Registry doesn't function at all in pure local-file mode (versions and stage transitions need a real database to coordinate concurrent writers safely).

Local File Store (Default)

import mlflow
# Metadata AND artifacts both land in ./mlruns on local disk.
# Fine for solo work; breaks down with concurrent writers or remote teammates.
mlflow.set_tracking_uri("file:./mlruns")

Centralized Tracking Server

A production MLflow setup separates three concerns, each backed by infrastructure suited to it:

# Launching a centralized MLflow tracking server
# (run this once, on a small persistent host/container — not on every laptop)
mlflow server \
  --backend-store-uri postgresql://mlflow_user:secret@db.internal:5432/mlflow_db \
  --default-artifact-root s3://my-company-mlflow-artifacts/ \
  --host 0.0.0.0 \
  --port 5000

# Every data scientist's laptop / training script now points at the
# shared server instead of a local folder — one line changes:
import mlflow

mlflow.set_tracking_uri("http://mlflow.internal.company.com:5000")
mlflow.set_experiment("churn-prediction")

# Everything else — log_param, log_metric, log_model, register_model —
# is IDENTICAL code to the local-file examples above. The tracking URI
# is the only thing that changes between solo and team mode.
with mlflow.start_run(run_name="rf-from-shared-server"):
    mlflow.log_params({"n_estimators": 300, "max_depth": 8})
    # ... train and log as before ...