MLflow Gradient Boosting Tutorial

Why This Matters

Gradient boosting is a practical choice for tabular machine learning, and MLflow adds the operational pieces that teams usually need next: experiment tracking, model registration, versioning, tags, and aliases for stable deployment targets such as champion and challenger. In open-source MLflow, the Model Registry is available through the UI and API, and alias-based loading lets inference code target a stable name instead of a hard-coded version number.

Step Summary

In this tutorial, you will:

Start a local MLflow server
Train a first gradient boosting model
Register it in the MLflow Model Registry
Tag and alias it as the champion
Run inference on the champion model
Train a second model as the challenger
Register and deploy the challenger locally
Run inference on the challenger deployment

Expected outcome: one registered model with at least two versions, a champion alias for the first version, a challenger alias for the second version, and working inference examples for both.

Requirements

Basic Python and scikit-learn experience
Python 3.10+
A local shell or terminal
MLflow, scikit-learn, pandas, and requests
No GPU required

1. Install the packages

We will use scikit-learn’s histogram-based gradient boosting classifier so the example stays compact and runs well on a laptop.

python -m venv .venv
source .venv/bin/activate
pip install "mlflow>=2.9" "scikit-learn>=1.6" pandas requests

If you are on Windows, activate the environment with the PowerShell activation script instead.

2. Start the MLflow server

For Model Registry features, MLflow should use a database-backed backend store instead of relying only on the legacy file backend. The server CLI supports a backend store and a separate default artifact root, which is exactly what we will configure here.

mkdir -p mlartifacts

mlflow server \
  --backend-store-uri sqlite:///mlflow.db \
  --default-artifact-root file:./mlartifacts \
  --host 127.0.0.1 \
  --port 5000

Open a second terminal and point your client code to this server:

export MLFLOW_TRACKING_URI=<your_mlflow_server_uri>

Keep the server terminal running during the rest of the tutorial.

3. Train and register the first GBM

In this step, we train a first gradient boosting classifier on the breast cancer dataset, log metrics, log the model artifact, and register it under one shared model name. We also include a model signature and input example, which MLflow supports directly during model logging. When registered_model_name is provided, MLflow creates a model version in the registry. (mlflow.org)

# train_champion.py
import os
import mlflow
from mlflow import MlflowClient
from mlflow.models import infer_signature
from sklearn.datasets import load_breast_cancer
from sklearn.ensemble import HistGradientBoostingClassifier
from sklearn.metrics import accuracy_score, roc_auc_score
from sklearn.model_selection import train_test_split

MODEL_NAME = "breast-cancer-gbm"
EXPERIMENT_NAME = "gbm-registry-demo"

mlflow.set_tracking_uri(os.environ["MLFLOW_TRACKING_URI"])
mlflow.set_experiment(EXPERIMENT_NAME)
client = MlflowClient()

def get_version_for_run(model_name, run_id):
    versions = client.search_model_versions(f"name='{model_name}'")
    matches = [v for v in versions if v.run_id == run_id]
    if not matches:
        raise RuntimeError(f"No registered version found for run_id={run_id}")
    return sorted(matches, key=lambda v: int(v.version))[-1].version

data = load_breast_cancer(as_frame=True)
X = data.data
y = data.target

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

params = {
    "learning_rate": 0.05,
    "max_depth": 3,
    "max_iter": 150,
    "min_samples_leaf": 20,
    "random_state": 42,
}

with mlflow.start_run(run_name="gbm-champion") as run:
    model = HistGradientBoostingClassifier(**params)
    model.fit(X_train, y_train)

    y_pred = model.predict(X_test)
    y_prob = model.predict_proba(X_test)[:, 1]

    mlflow.log_params(params)
    mlflow.log_metric("accuracy", accuracy_score(y_test, y_pred))
    mlflow.log_metric("roc_auc", roc_auc_score(y_test, y_prob))

    signature = infer_signature(X_train, model.predict(X_train))

    mlflow.sklearn.log_model(
        sk_model=model,
        artifact_path="model",
        input_example=X_train.head(3),
        signature=signature,
        registered_model_name=MODEL_NAME,
    )

version = get_version_for_run(MODEL_NAME, run.info.run_id)

client.set_model_version_tag(MODEL_NAME, version, "role", "champion")
client.set_registered_model_alias(MODEL_NAME, "champion", version)

print(f"Registered champion version: {version}")

Run it:

python train_champion.py

4. Confirm the champion tag and alias

We use both a tag and an alias here. The tag stores metadata such as role=champion, while the alias gives us a stable reference we can load later with a model URI like models:/<model-name>@champion. That means inference code can stay the same even if the alias is moved to a newer version later.

# inspect_registry.py
import os
import mlflow
from mlflow import MlflowClient

MODEL_NAME = "breast-cancer-gbm"

mlflow.set_tracking_uri(os.environ["MLFLOW_TRACKING_URI"])
client = MlflowClient()

champion = client.get_model_version_by_alias(MODEL_NAME, "champion")

print("Champion version:", champion.version)
print("Champion tags:", dict(champion.tags))

Run it:

python inspect_registry.py

5. Run inference on the champion model

Now we load the registered model by alias instead of by a hard-coded version number. This is the simplest pattern for application code because the alias stays stable while the target version can change behind the scenes.

# infer_champion.py
import os
import pandas as pd
import mlflow
from sklearn.datasets import load_breast_cancer

MODEL_NAME = "breast-cancer-gbm"

mlflow.set_tracking_uri(os.environ["MLFLOW_TRACKING_URI"])

model = mlflow.sklearn.load_model(f"models:/{MODEL_NAME}@champion")

X = load_breast_cancer(as_frame=True).data.head(5)

predictions = pd.DataFrame(
    {
        "prediction": model.predict(X),
        "prob_positive": model.predict_proba(X)[:, 1],
    }
)

print(predictions)

Run it:

python infer_champion.py

6. Train and register the challenger model

Next, we train a second model with different hyperparameters. It is registered under the same model name, which creates another version. Then we tag and alias it as challenger. Model Registry versions are meant for exactly this kind of side-by-side candidate management.

# train_challenger.py
import os
import mlflow
from mlflow import MlflowClient
from mlflow.models import infer_signature
from sklearn.datasets import load_breast_cancer
from sklearn.ensemble import HistGradientBoostingClassifier
from sklearn.metrics import accuracy_score, roc_auc_score
from sklearn.model_selection import train_test_split

MODEL_NAME = "breast-cancer-gbm"
EXPERIMENT_NAME = "gbm-registry-demo"

mlflow.set_tracking_uri(os.environ["MLFLOW_TRACKING_URI"])
mlflow.set_experiment(EXPERIMENT_NAME)
client = MlflowClient()

def get_version_for_run(model_name, run_id):
    versions = client.search_model_versions(f"name='{model_name}'")
    matches = [v for v in versions if v.run_id == run_id]
    if not matches:
        raise RuntimeError(f"No registered version found for run_id={run_id}")
    return sorted(matches, key=lambda v: int(v.version))[-1].version

data = load_breast_cancer(as_frame=True)
X = data.data
y = data.target

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

params = {
    "learning_rate": 0.10,
    "max_depth": 5,
    "max_iter": 250,
    "min_samples_leaf": 10,
    "random_state": 42,
}

with mlflow.start_run(run_name="gbm-challenger") as run:
    model = HistGradientBoostingClassifier(**params)
    model.fit(X_train, y_train)

    y_pred = model.predict(X_test)
    y_prob = model.predict_proba(X_test)[:, 1]

    mlflow.log_params(params)
    mlflow.log_metric("accuracy", accuracy_score(y_test, y_pred))
    mlflow.log_metric("roc_auc", roc_auc_score(y_test, y_prob))

    signature = infer_signature(X_train, model.predict(X_train))

    mlflow.sklearn.log_model(
        sk_model=model,
        artifact_path="model",
        input_example=X_train.head(3),
        signature=signature,
        registered_model_name=MODEL_NAME,
    )

version = get_version_for_run(MODEL_NAME, run.info.run_id)

client.set_model_version_tag(MODEL_NAME, version, "role", "challenger")
client.set_registered_model_alias(MODEL_NAME, "challenger", version)

print(f"Registered challenger version: {version}")

Run it:

python train_challenger.py

You can update the inspection script to print both aliases:

# inspect_both.py
import os
import mlflow
from mlflow import MlflowClient

MODEL_NAME = "breast-cancer-gbm"

mlflow.set_tracking_uri(os.environ["MLFLOW_TRACKING_URI"])
client = MlflowClient()

for alias in ["champion", "challenger"]:
    mv = client.get_model_version_by_alias(MODEL_NAME, alias)
    print(alias, "-> version", mv.version, "tags:", dict(mv.tags))

7. Deploy the challenger model locally

MLflow can serve a registered model as a local inference server with mlflow models serve. Also note that the CLI should be pointed at your tracking server with MLFLOW_TRACKING_URI; otherwise it falls back to the local filesystem instead of your running MLflow server.

export MLFLOW_TRACKING_URI=<your_mlflow_server_uri>

mlflow models serve \
  -m "models:/breast-cancer-gbm@challenger" \
  -p 5001 \
  --env-manager local

Leave this serving process running.

8. Run inference on the challenger deployment

The local serving endpoint accepts structured prediction payloads such as dataframe_split. Here we send three rows from the same dataset to the challenger server. The endpoint is read from an environment variable so you can reuse the script in any local setup.

# call_challenger.py
import os
import json
import requests
from sklearn.datasets import load_breast_cancer

endpoint = os.environ["CHALLENGER_INVOCATIONS_ENDPOINT"]

X = load_breast_cancer(as_frame=True).data.head(3)
payload = {"dataframe_split": X.to_dict(orient="split")}

response = requests.post(
    endpoint,
    headers={"Content-Type": "application/json"},
    data=json.dumps(payload),
    timeout=30,
)

print("Status:", response.status_code)
print("Body:", response.json())

Set the endpoint and call the service:

export CHALLENGER_INVOCATIONS_ENDPOINT=<your_challenger_invocations_endpoint>
python call_challenger.py

Recap

You now have a full small-scale MLOps workflow:

a first gradient boosting model trained and registered
a champion tag and alias attached to that version
inference running against the champion alias
a second challenger model registered as a new version
a local deployment of the challenger model
online inference against the challenger endpoint

This pattern scales well because your application can load by alias, while the registry keeps track of versions, tags, and lineage.

FAQ

1. Why use both a tag and an alias?

A tag is metadata. It helps you describe a model version, such as role=champion. An alias is a movable pointer that your inference code can load directly. In practice, teams often keep both.

2. Why did my version numbers start at 3 or 4 instead of 1 or 2?

MLflow versions increment inside the registered model. If you rerun the tutorial with the same model name, new versions are appended. Use a new model name if you want a clean sequence.

3. Do I need a GPU for this tutorial?

No. The entire example runs on CPU and is intentionally small enough for a normal laptop or workstation.

MLflow Gradient Boosting Tutorial

Why This Matters

Step Summary

Requirements

1. Install the packages

2. Start the MLflow server

3. Train and register the first GBM

4. Confirm the champion tag and alias

5. Run inference on the champion model

6. Train and register the challenger model

7. Deploy the challenger model locally

8. Run inference on the challenger deployment

Recap

Further Reading

FAQ

1. Why use both a tag and an alias?

2. Why did my version numbers start at 3 or 4 instead of 1 or 2?

3. Do I need a GPU for this tutorial?

Leave a Comment Cancel Reply

Why This Matters

Step Summary

Requirements

1. Install the packages

2. Start the MLflow server

3. Train and register the first GBM

4. Confirm the champion tag and alias

5. Run inference on the champion model

6. Train and register the challenger model

7. Deploy the challenger model locally

8. Run inference on the challenger deployment

Recap

Further Reading

FAQ

1. Why use both a tag and an alias?

2. Why did my version numbers start at 3 or 4 instead of 1 or 2?

3. Do I need a GPU for this tutorial?

Related Posts

Leave a Comment Cancel Reply