XGBoost¶

XGBoost is an optimized distributed gradient boosting library designed to be highly efficient, flexible and portable. It implements machine learning algorithms under the Gradient Boosting framework.

This document explains how to serve and deploy an XGBoost model for predicting breast cancer with BentoML.

You can query the exposed prediction endpoint with breast tumor data. For example:

{
     "data": [
                 [1.308e+01, 1.571e+01, 8.563e+01, 5.200e+02, 1.075e-01, 1.270e-01,
                 4.568e-02, 3.110e-02, 1.967e-01, 6.811e-02, 1.852e-01, 7.477e-01,
                 1.383e+00, 1.467e+01, 4.097e-03, 1.898e-02, 1.698e-02, 6.490e-03,
                 1.678e-02, 2.425e-03, 1.450e+01, 2.049e+01, 9.609e+01, 6.305e+02,
                 1.312e-01, 2.776e-01, 1.890e-01, 7.283e-02, 3.184e-01, 8.183e-02]
             ]
 }

Expected output:

[[0.02664177 0.9733583 ]] # 2.66% chance benign and 97.34% chance malignant

Code explanations¶

You can find the source code in GitHub. Below is a breakdown of the key code implementations within this project.

save_model.py¶

This example uses the scikit-learn framework to load and preprocess the breast cancer dataset, which is then converted into an XGBoost-compatible format (DMatrix) to train the machine learning model.

save_model.py¶
import typing as t
from sklearn.datasets import load_breast_cancer
from sklearn.utils import Bunch
import xgboost as xgb

# Load the data
cancer: Bunch = t.cast("Bunch", load_breast_cancer())
cancer_data = t.cast("ext.NpNDArray", cancer.data)
cancer_target = t.cast("ext.NpNDArray", cancer.target)
dt = xgb.DMatrix(cancer_data, label=cancer_target)

# Specify model parameters
param = {
    "max_depth": 3,
    "eta": 0.3,
    "objective": "multi:softprob",
    "num_class": 2
}

# Train the model
model = xgb.train(param, dt)

After training, use the bentoml.xgboost.save_model API to save the model to the BentoML Model Store, a local directory to store and manage models. You can retrieve this model later in other services to run predictions.

import bentoml

# Specify the model name and the model to be saved
bentoml.xgboost.save_model("cancer", model)

To verify that the model has been successfully saved, run:

$ bentoml models list

Tag                      Module           Size       Creation Time
cancer:xa2npbboccvv7u4c  bentoml.xgboost  23.17 KiB  2024-06-19 07:51:21

test.py¶

To ensure that the saved model works correctly, try loading it and running a prediction:

test.py¶
import bentoml
import xgboost as xgb

# Load the model by setting the model tag
booster = bentoml.xgboost.load_model("cancer:xa2npbboccvv7u4c")

# Predict using a sample
res = booster.predict(xgb.DMatrix([[1.308e+01, 1.571e+01, 8.563e+01, 5.200e+02, 1.075e-01, 1.270e-01,
    4.568e-02, 3.110e-02, 1.967e-01, 6.811e-02, 1.852e-01, 7.477e-01,
    1.383e+00, 1.467e+01, 4.097e-03, 1.898e-02, 1.698e-02, 6.490e-03,
    1.678e-02, 2.425e-03, 1.450e+01, 2.049e+01, 9.609e+01, 6.305e+02,
    1.312e-01, 2.776e-01, 1.890e-01, 7.283e-02, 3.184e-01, 8.183e-02]]))

print(res)

Expected result:

[[0.02664177 0.9733583 ]]

service.py¶

The service.py file is where you define the serving logic and expose the model as a web service.

service.py¶
import bentoml
import numpy as np
import xgboost as xgb
import os

@bentoml.service(
    resources={"cpu": "2"},
    traffic={"timeout": 10},
)
class CancerClassifier:
    # Declare the model as a class variable
    bento_model = bentoml.models.BentoModel("cancer:latest")

    def __init__(self):
        self.model = bentoml.xgboost.load_model(self.bento_model)

        # Check resource availability
        if os.getenv("CUDA_VISIBLE_DEVICES") not in (None, "", "-1"):
            self.model.set_param({"predictor": "gpu_predictor", "gpu_id": 0})  # type: ignore (incomplete XGBoost types)
        else:
            nthreads = os.getenv("OMP_NUM_THREADS")
            if nthreads:
                nthreads = max(int(nthreads), 1)
            else:
                nthreads = 1
            self.model.set_param(
                {"predictor": "cpu_predictor", "nthread": nthreads}
            )

    @bentoml.api
    def predict(self, data: np.ndarray) -> np.ndarray:
        return self.model.predict(xgb.DMatrix(data))

The Service code:

  • Uses the @bentoml.service decorator to define a BentoML Service. Optionally, you can set additional configurations like traffic timeout.

  • Retrieves the model from the Model Store and defines it a class variable.

  • Checks resource availability like GPUs and the number of threads.

  • Uses the @bentoml.api decorator to expose the predict function as an API endpoint, which takes a NumPy array as input and returns a NumPy array. Note that the input data is converted into a DMatrix, which is the data structure XGBoost uses for datasets.

The @bentoml.service decorator also allows you to define the runtime environment for a Bento, the unified distribution format in BentoML. A Bento is packaged with all the source code, Python dependencies, model references, and environment setup, making it easy to deploy consistently across different environments.

Here is an example:

service.py¶
my_image = bentoml.images.Image(python_version="3.11") \
            .python_packages("xgboost", "scikit-learn")

@bentoml.service(
    image=my_image, # Apply the specifications
    ...
)
class CancerClassifier:
    ...

Try it out¶

You can serve this example project locally, or containerize it as an OCI-compliant image and deploy it anywhere.

  1. Clone the project repository and install the dependencies.

    git clone https://github.com/bentoml/BentoXGBoost.git
cd BentoXGBoost

# Recommend Python 3.11
pip install bentoml xgboost scikit-learn

  2. Train and save the model to the BentoML Model Store.

    python3 save_model.py

  3. Serve it locally.

    bentoml serve

  4. Call the endpoint at http://localhost:3000. You can also open this URL in your browser to try it out in the Swagger UI.

    Create a BentoML client to call the endpoint.

    import bentoml

with bentoml.SyncHTTPClient("http://localhost:3000") as client:
    result = client.predict(
        data=[
            [1.308e+01, 1.571e+01, 8.563e+01, 5.200e+02, 1.075e-01, 1.270e-01,
             4.568e-02, 3.110e-02, 1.967e-01, 6.811e-02, 1.852e-01, 7.477e-01,
             1.383e+00, 1.467e+01, 4.097e-03, 1.898e-02, 1.698e-02, 6.490e-03,
             1.678e-02, 2.425e-03, 1.450e+01, 2.049e+01, 9.609e+01, 6.305e+02,
             1.312e-01, 2.776e-01, 1.890e-01, 7.283e-02, 3.184e-01, 8.183e-02]
        ],
    )
    print(result)

    curl -X POST http://localhost:3000/predict \
    -H 'Content-Type: application/json' \
    -d '{
        "data": [
            [1.308e+01, 1.571e+01, 8.563e+01, 5.200e+02, 1.075e-01, 1.270e-01,
             4.568e-02, 3.110e-02, 1.967e-01, 6.811e-02, 1.852e-01, 7.477e-01,
             1.383e+00, 1.467e+01, 4.097e-03, 1.898e-02, 1.698e-02, 6.490e-03,
             1.678e-02, 2.425e-03, 1.450e+01, 2.049e+01, 9.609e+01, 6.305e+02,
             1.312e-01, 2.776e-01, 1.890e-01, 7.283e-02, 3.184e-01, 8.183e-02]
        ]
    }'

For custom deployment in your own infrastructure, use BentoML to generate an OCI-compliant image.