Repository: bentoml/BentoML

Issue:
feature: add support for XGBoost models trained using the XGBoost scikit-learn interface
### Feature request

I would like to use BentoML to save and deploy XGBoost models that were trained using the scikit-learn learn interface of the XGBoost model. 



### Motivation

I'm trying to save an XGBoost model using BentoML. I'm using the [scikit-learn interface](https://xgboost.readthedocs.io/en/stable/python/python_api.html#module-xgboost.sklearn) of the XGBoost module, so my final model is of type `xgboost.sklearn.XGBClassifier`. I noticed [here](https://github.com/bentoml/BentoML/blob/main/bentoml/_internal/frameworks/xgboost.py#L162) that BentoML expects the model to be of type `xgboost.core.Booster`, which is what you get when you use the native XGBoost module. 

I also can't save the model as an sklearn model (without failing [this validation](https://github.com/bentoml/BentoML/blob/main/bentoml/_internal/frameworks/sklearn.py#L145)).

### Other

_No response_


Discussion:
Hi @gustavocidornelas, I believe the support for sklearn interface are straightforward, I will work on this and get back to you asap.

Resolving patch:
diff --git a/bentoml/_internal/frameworks/xgboost.py b/bentoml/_internal/frameworks/xgboost.py
--- a/bentoml/_internal/frameworks/xgboost.py
+++ b/bentoml/_internal/frameworks/xgboost.py
@@ -6,12 +6,15 @@
 from types import ModuleType
 from typing import TYPE_CHECKING
 
+import attr
 import numpy as np
 
 import bentoml
 from bentoml import Tag
+from bentoml.models import ModelOptions
 from bentoml.exceptions import NotFound
 from bentoml.exceptions import InvalidArgument
+from bentoml.exceptions import BentoMLException
 from bentoml.exceptions import MissingDependencyException
 from bentoml._internal.models.model import ModelContext
 
@@ -32,13 +35,27 @@
         "https://xgboost.readthedocs.io/en/latest/install.html"
     )
 
+try:
+    from xgboost import XGBModel
+except ImportError:  # pragma: no cover
+    # if sklearn is not installed, XGBoost will not expose XGBModel, so make
+    # a dummy class ourselves
+    class XGBModel:
+        pass
+
+
 MODULE_NAME = "bentoml.xgboost"
 MODEL_FILENAME = "saved_model.ubj"
-API_VERSION = "v1"
+API_VERSION = "v2"
 
 logger = logging.getLogger(__name__)
 
 
+@attr.define
+class XGBoostOptions(ModelOptions):
+    model_class: str | None = None
+
+
 def get(tag_like: str | Tag) -> bentoml.Model:
     """
     Get the BentoML model with the given tag.
@@ -64,7 +81,9 @@ def get(tag_like: str | Tag) -> bentoml.Model:
     return model
 
 
-def load_model(bento_model: str | Tag | bentoml.Model) -> xgb.core.Booster:
+def load_model(
+    bento_model: str | Tag | bentoml.Model,
+) -> xgb.Booster | xgb.XGBModel:
     """
     Load the XGBoost model with the given tag from the local BentoML model store.
 
@@ -73,7 +92,7 @@ def load_model(bento_model: str | Tag | bentoml.Model) -> xgb.core.Booster:
             Either the tag of the model to get from the store, or a BentoML `~bentoml.Model`
             instance to load the model from.
     Returns:
-        :obj:`~xgboost.core.Booster`: The XGBoost model loaded from the model store or BentoML :obj:`~bentoml.Model`.
+        The XGBoost model loaded from the model store or BentoML :obj:`~bentoml.Model`.
     Example:
 
     .. code-block:: python
@@ -92,13 +111,40 @@ def load_model(bento_model: str | Tag | bentoml.Model) -> xgb.core.Booster:
         )
 
     model_file = bento_model.path_of(MODEL_FILENAME)
-    booster = xgb.core.Booster(model_file=model_file)
-    return booster
+    model_api_version = bento_model.info.api_version
+    if model_api_version == "v1":
+        model = xgb.Booster(model_file=model_file)
+    else:
+        if model_api_version != "v2":
+            logger.warning(
+                f"Got an XGBoost model with an unsupported version '{model_api_version}, unexpected errors may occur."
+            )
+        model_class = t.cast(XGBoostOptions, bento_model.info.options).model_class
+        if model_class is None:
+            raise BentoMLException(
+                f"Model '{bento_model.tag}' is missing the required 'model_class' option. This should not be possible; please file an issue if you encounter this error."
+            )
+        try:
+            xgb_class: type[xgb.XGBModel] | type[xgb.Booster] = getattr(
+                xgb, model_class
+            )
+        except AttributeError:
+            if model_class != "Booster":
+                raise BentoMLException(
+                    f"Model '{bento_model.tag}' is an XGBoost Scikit-Learn model, but sklearn is not installed."
+                ) from None
+            else:
+                raise BentoMLException(
+                    "xgboost.Booster could not be found, your XGBoost installation may be corrupted. Ensure there is no file named 'xgboost.py' that may be being loaded instead of the XGBoost library."
+                ) from None
+        model: xgb.Booster | xgb.XGBModel = xgb_class()
+        model.load_model(model_file)
+    return model
 
 
 def save_model(
     name: str,
-    model: xgb.core.Booster,
+    model: xgb.Booster | xgb.XGBModel,
     *,
     signatures: dict[str, ModelSignatureDict] | None = None,
     labels: dict[str, str] | None = None,
@@ -110,35 +156,34 @@ def save_model(
     Save an XGBoost model instance to the BentoML model store.
 
     Args:
-        name (``str``):
+        name:
             The name to give to the model in the BentoML store. This must be a valid
             :obj:`~bentoml.Tag` name.
-        model (:obj:`~xgboost.core.Booster`):
+        model:
             The XGBoost model to be saved.
-        signatures (``dict[str, ModelSignatureDict]``, optional):
+        signatures:
             Signatures of predict methods to be used. If not provided, the signatures default to
             ``{"predict": {"batchable": False}}``. See :obj:`~bentoml.types.ModelSignature` for more
             details.
-        labels (``dict[str, str]``, optional):
+        labels:
             A default set of management labels to be associated with the model. An example is
             ``{"training-set": "data-1"}``.
-        custom_objects (``dict[str, Any]``, optional):
+        custom_objects:
             Custom objects to be saved with the model. An example is
             ``{"my-normalizer": normalizer}``.
 
             Custom objects are currently serialized with cloudpickle, but this implementation is
             subject to change.
-        external_modules (:code:`List[ModuleType]`, `optional`, default to :code:`None`):
+        external_modules:
             user-defined additional python modules to be saved alongside the model or custom objects,
             e.g. a tokenizer module, preprocessor module, model configuration module
-        metadata (``dict[str, Any]``, optional):
+        metadata:
             Metadata to be associated with the model. An example is ``{"max_depth": 2}``.
 
             Metadata is intended for display in model management UI and therefore must be a default
             Python type, such as ``str`` or ``int``.
     Returns:
-        :obj:`~bentoml.Tag`: A :obj:`tag` with a format `name:version` where `name` is the
-        user-defined model's name, and a generated `version` by BentoML.
+        A BentoML tag with the user-defined name and a generated version.
 
     Example:
 
@@ -159,8 +204,12 @@ def save_model(
         # `save` the booster to BentoML modelstore:
         bento_model = bentoml.xgboost.save_model("my_xgboost_model", bst, booster_params=param)
     """  # noqa: LN001
-    if not isinstance(model, xgb.core.Booster):
-        raise TypeError(f"Given model ({model}) is not a xgboost.core.Booster.")
+    if isinstance(model, xgb.Booster):
+        model_class = "Booster"
+    elif isinstance(model, XGBModel):
+        model_class = model.__class__.__name__
+    else:
+        raise TypeError(f"Given model ({model}) is not a xgboost.Booster.")
 
     context: ModelContext = ModelContext(
         framework_name="xgboost",
@@ -185,6 +234,7 @@ def save_model(
         external_modules=external_modules,
         metadata=metadata,
         context=context,
+        options=XGBoostOptions(model_class=model_class),
     ) as bento_model:
         model.save_model(bento_model.path_of(MODEL_FILENAME))  # type: ignore (incomplete XGBoost types)
 
@@ -204,17 +254,22 @@ def __init__(self):
             super().__init__()
             self.model = load_model(bento_model)
 
+            self.booster = (
+                self.model
+                if isinstance(self.model, xgb.Booster)
+                else self.model.get_booster()
+            )
+
             # check for resources
-            available_gpus = os.getenv("CUDA_VISIBLE_DEVICES")
-            if available_gpus is not None and available_gpus not in ("", "-1"):
-                self.model.set_param({"predictor": "gpu_predictor", "gpu_id": 0})  # type: ignore (incomplete XGBoost types)
+            if os.getenv("CUDA_VISIBLE_DEVICES") not in (None, "", "-1"):
+                self.booster.set_param({"predictor": "gpu_predictor", "gpu_id": 0})  # type: ignore (incomplete XGBoost types)
             else:
                 nthreads = os.getenv("OMP_NUM_THREADS")
                 if nthreads is not None and nthreads != "":
                     nthreads = max(int(nthreads), 1)
                 else:
                     nthreads = 1
-                self.model.set_param({"predictor": "cpu_predictor", "nthread": nthreads})  # type: ignore (incomplete XGBoost types)
+                self.booster.set_param({"predictor": "cpu_predictor", "nthread": nthreads})  # type: ignore (incomplete XGBoost types)
 
             self.predict_fns: dict[str, t.Callable[..., t.Any]] = {}
             for method_name in bento_model.info.signatures:
@@ -231,9 +286,12 @@ def _run(
             input_data: ext.NpNDArray
             | ext.PdDataFrame,  # TODO: add support for DMatrix
         ) -> ext.NpNDArray:
-            dmatrix = xgb.DMatrix(input_data)
+            if isinstance(self.model, xgb.Booster):
+                inp = xgb.DMatrix(input_data)
+            else:
+                inp = input_data
 
-            res = self.predict_fns[method_name](dmatrix)
+            res = self.predict_fns[method_name](inp)
             return np.asarray(res)  # type: ignore (incomplete np types)
 
         XGBoostRunnable.add_method(
diff --git a/bentoml/xgboost.py b/bentoml/xgboost.py
--- a/bentoml/xgboost.py
+++ b/bentoml/xgboost.py
@@ -4,6 +4,7 @@
 from ._internal.frameworks.xgboost import load_model
 from ._internal.frameworks.xgboost import save_model
 from ._internal.frameworks.xgboost import get_runnable
+from ._internal.frameworks.xgboost import XGBoostOptions as ModelOptions
 
 logger = logging.getLogger(__name__)
 
@@ -34,4 +35,4 @@ def load_runner(tag, *args, **kwargs):
     return get(tag).to_runner()
 
 
-__all__ = ["load_model", "save_model", "get", "get_runnable"]
+__all__ = ["load_model", "save_model", "get", "get_runnable", "ModelOptions"]


Test patch:
diff --git a/tests/integration/frameworks/models/__init__.py b/tests/integration/frameworks/models/__init__.py
--- a/tests/integration/frameworks/models/__init__.py
+++ b/tests/integration/frameworks/models/__init__.py
@@ -3,6 +3,7 @@
 import typing as t
 
 import attr
+import numpy as np
 
 
 @attr.define
@@ -47,6 +48,9 @@ def check_output(self, outp: t.Any):
                     result
                 ), f"Output from model call ({', '.join(map(str, self.input_args))}, **{self.input_kwargs}) is not as expected"
         else:
+            check = outp == self.expected
+            if isinstance(check, np.ndarray):
+                check = check.all()
             assert (
-                outp == self.expected
+                check
             ), f"Output from model call ({', '.join(map(str, self.input_args))}, **{self.input_kwargs}) is not as expected"
diff --git a/tests/integration/frameworks/models/xgboost.py b/tests/integration/frameworks/models/xgboost.py
--- a/tests/integration/frameworks/models/xgboost.py
+++ b/tests/integration/frameworks/models/xgboost.py
@@ -6,6 +6,7 @@
 import numpy as np
 import pandas as pd
 import xgboost as xgb
+from sklearn.datasets import load_iris
 from sklearn.datasets import load_breast_cancer
 
 import bentoml
@@ -24,7 +25,9 @@
 param = {"max_depth": 3, "eta": 0.3, "objective": "multi:softprob", "num_class": 2}
 
 
-def check_model(bst: xgb.Booster, resources: dict[str, t.Any]):
+def check_model(bst: xgb.Booster | xgb.XGBModel, resources: dict[str, t.Any]):
+    if not isinstance(bst, xgb.Booster):
+        bst = bst.get_booster()
     config = json.loads(bst.save_config())
     if "nvidia.com/gpu" in resources and resources["nvidia.com/gpu"] > 0:
         assert config["learner"]["generic_param"]["nthread"] == str(1)
@@ -81,4 +84,50 @@ def check_output(out):
         ),
     ],
 )
-models: list[FrameworkTestModel] = [cancer_model]
+
+
+iris = load_iris()
+
+xgbc = xgb.XGBClassifier()
+xgbc.fit(iris.data, iris.target)
+
+
+iris_model = FrameworkTestModel(
+    name="iris",
+    model=xgbc,
+    configurations=[
+        Config(
+            test_inputs={
+                "predict": [
+                    Input(
+                        input_args=[np.array([iris.data[0]])],
+                        expected=[iris.target[0]],
+                    ),
+                    Input(
+                        input_args=[np.array(iris.data[1:5])],
+                        expected=iris.target[1:5],
+                    ),
+                ],
+            },
+            check_model=check_model,
+        ),
+        Config(
+            test_inputs={
+                "predict": [
+                    Input(
+                        input_args=[pd.DataFrame([iris.data[0]])],
+                        expected=[iris.target[0]],
+                    ),
+                    Input(
+                        input_args=[pd.DataFrame(iris.data[1:5])],
+                        expected=iris.target[1:5],
+                    ),
+                ],
+            },
+            check_model=check_model,
+        ),
+    ],
+)
+
+
+models: list[FrameworkTestModel] = [cancer_model, iris_model]