Repository: SeldonIO/MLServer

Issue:
Add support for predict_proba in XGBoost inference runtime
Currently, the XGBoost runtime only supports the `predict()` method. On top of this one, it would be useful to also add support for the [`predict_proba()` one](https://xgboost.readthedocs.io/en/stable/python/python_api.html#xgboost.XGBClassifier.predict_proba), similar to how the SKLearn runtime works.


Discussion:
The `xgboost.Booster.predict` method used in the XGBoost runtime already returns values similar to `predict_proba` (i.e. probabilities and not predictions).

If the user wants to use the `xgboost.sklearn.XGBClassifier` interface, they can always save their `XGBClassifier` with joblib and use the "mlserver_sklearn.SKLearnModel` interface.

Resolving patch:
diff --git a/runtimes/xgboost/mlserver_xgboost/codecs.py b/runtimes/xgboost/mlserver_xgboost/codecs.py
deleted file mode 100644
--- a/runtimes/xgboost/mlserver_xgboost/codecs.py
+++ /dev/null
@@ -1,33 +0,0 @@
-from xgboost import DMatrix
-
-from typing import Any
-
-from mlserver.codecs import NumpyCodec, register_input_codec, register_request_codec
-from mlserver.codecs.utils import SingleInputRequestCodec
-from mlserver.types import RequestInput
-from mlserver.errors import InferenceError
-
-
-@register_input_codec
-class DMatrixCodec(NumpyCodec):
-    ContentType = "dmatrix"
-
-    @classmethod
-    def can_encode(cls, payload: Any) -> bool:
-        return isinstance(payload, DMatrix)
-
-    @classmethod
-    def decode_input(cls, request_input: RequestInput) -> DMatrix:  # type: ignore
-        try:
-            ndarray = super().decode_input(request_input)
-            return DMatrix(ndarray)
-        except Exception as e:
-            # There are a few things that can go wrong here, e.g. less than 2-D
-            # in the array), or input data not compatible with a numpy array
-            raise InferenceError("Invalid input to XGBoostModel") from e
-
-
-@register_request_codec
-class DMatrixRequestCodec(SingleInputRequestCodec):
-    InputCodec = DMatrixCodec
-    ContentType = DMatrixCodec.ContentType
diff --git a/runtimes/xgboost/mlserver_xgboost/xgboost.py b/runtimes/xgboost/mlserver_xgboost/xgboost.py
--- a/runtimes/xgboost/mlserver_xgboost/xgboost.py
+++ b/runtimes/xgboost/mlserver_xgboost/xgboost.py
@@ -1,36 +1,95 @@
 import xgboost as xgb
 
-from mlserver import types
+from typing import List
+from xgboost.sklearn import XGBModel
+
+from mlserver.errors import InferenceError
 from mlserver.model import MLModel
 from mlserver.utils import get_model_uri
+from mlserver.codecs import NumpyRequestCodec, NumpyCodec
+from mlserver.types import (
+    InferenceRequest,
+    InferenceResponse,
+    RequestOutput,
+    ResponseOutput,
+)
 
-from .codecs import DMatrixCodec, DMatrixRequestCodec
-
+PREDICT_OUTPUT = "predict"
+PREDICT_PROBA_OUTPUT = "predict_proba"
+VALID_OUTPUTS = [PREDICT_OUTPUT, PREDICT_PROBA_OUTPUT]
 
 WELLKNOWN_MODEL_FILENAMES = ["model.bst", "model.json"]
 
 
+def _load_sklearn_interface(model_uri: str) -> XGBModel:
+    try:
+        regressor = xgb.XGBRegressor()
+        regressor.load_model(model_uri)
+        return regressor
+    except TypeError:
+        # If there was an error, it's likely due to the model being a
+        # classifier
+        classifier = xgb.XGBClassifier()
+        classifier.load_model(model_uri)
+        return classifier
+
+
 class XGBoostModel(MLModel):
     """
     Implementationof the MLModel interface to load and serve `xgboost` models.
     """
 
     async def load(self) -> bool:
-        # TODO: Log info message
         model_uri = await get_model_uri(
             self._settings, wellknown_filenames=WELLKNOWN_MODEL_FILENAMES
         )
-        self._model = xgb.Booster(model_file=model_uri)
+
+        self._model = _load_sklearn_interface(model_uri)
 
         self.ready = True
         return self.ready
 
-    async def predict(self, payload: types.InferenceRequest) -> types.InferenceResponse:
-        decoded = self.decode_request(payload, default_codec=DMatrixRequestCodec)
-        prediction = self._model.predict(decoded)
+    def _check_request(self, payload: InferenceRequest) -> InferenceRequest:
+        if not payload.outputs:
+            # By default, only return the result of `predict()`
+            payload.outputs = [RequestOutput(name=PREDICT_OUTPUT)]
+        else:
+            for request_output in payload.outputs:
+                if request_output.name not in VALID_OUTPUTS:
+                    raise InferenceError(
+                        f"XGBoostModel only supports '{PREDICT_OUTPUT}' and "
+                        f"'{PREDICT_PROBA_OUTPUT}' as outputs "
+                        f"({request_output.name} was received)"
+                    )
+
+        # Regression models do not support `predict_proba`
+        if PREDICT_PROBA_OUTPUT in [o.name for o in payload.outputs]:
+            if isinstance(self._model, xgb.XGBRegressor):
+                raise InferenceError(
+                    f"XGBRegressor models do not support '{PREDICT_PROBA_OUTPUT}"
+                )
+
+        return payload
+
+    def _get_model_outputs(self, payload: InferenceRequest) -> List[ResponseOutput]:
+        decoded_request = self.decode_request(payload, default_codec=NumpyRequestCodec)
+
+        outputs = []
+        for request_output in payload.outputs:  # type: ignore
+            predict_fn = getattr(self._model, request_output.name)
+            y = predict_fn(decoded_request)
+
+            output = self.encode(y, request_output, default_codec=NumpyCodec)
+            outputs.append(output)
+
+        return outputs
+
+    async def predict(self, payload: InferenceRequest) -> InferenceResponse:
+        payload = self._check_request(payload)
+        outputs = self._get_model_outputs(payload)
 
-        return types.InferenceResponse(
+        return InferenceResponse(
             model_name=self.name,
             model_version=self.version,
-            outputs=[DMatrixCodec.encode_output("predict", prediction)],
+            outputs=outputs,
         )
diff --git a/runtimes/xgboost/setup.py b/runtimes/xgboost/setup.py
--- a/runtimes/xgboost/setup.py
+++ b/runtimes/xgboost/setup.py
@@ -33,7 +33,7 @@ def _load_description() -> str:
     author_email="hello@seldon.io",
     description="XGBoost runtime for MLServer",
     packages=find_packages(exclude=["tests", "tests.*"]),
-    install_requires=["mlserver", "xgboost"],
+    install_requires=["mlserver", "xgboost[scikit-learn]"],
     long_description=_load_description(),
     long_description_content_type="text/markdown",
     license="Apache 2.0",


Test patch:
diff --git a/runtimes/xgboost/tests/conftest.py b/runtimes/xgboost/tests/conftest.py
--- a/runtimes/xgboost/tests/conftest.py
+++ b/runtimes/xgboost/tests/conftest.py
@@ -37,6 +37,25 @@ def model_uri(tmp_path) -> str:
     return model_uri
 
 
+@pytest.fixture
+def classifier_uri(tmp_path) -> str:
+    n = 5
+    d = 3
+    c = 5
+
+    X_train = np.random.rand(n, d)
+    y_train = [label for label in range(0, c)]
+    clf = xgb.XGBClassifier(
+        num_class=c, use_label_encoder=False, objective="multi:softprob"
+    )
+    clf.fit(X_train, y_train, eval_metric="mlogloss")
+
+    model_uri = os.path.join(tmp_path, "xgboost-model.json")
+    clf.save_model(model_uri)
+
+    return model_uri
+
+
 @pytest.fixture
 def model_settings(model_uri: str) -> ModelSettings:
     return ModelSettings(
@@ -53,6 +72,18 @@ async def model(model_settings: ModelSettings) -> XGBoostModel:
     return model
 
 
+@pytest.fixture
+async def classifier(classifier_uri: str) -> XGBoostModel:
+    model_settings = ModelSettings(
+        name="xgboost-model",
+        parameters=ModelParameters(uri=classifier_uri, version="v1.2.3"),
+    )
+    model = XGBoostModel(model_settings)
+    await model.load()
+
+    return model
+
+
 @pytest.fixture
 def inference_request() -> InferenceRequest:
     payload_path = os.path.join(TESTDATA_PATH, "inference-request.json")
diff --git a/runtimes/xgboost/tests/test_xgboost.py b/runtimes/xgboost/tests/test_xgboost.py
--- a/runtimes/xgboost/tests/test_xgboost.py
+++ b/runtimes/xgboost/tests/test_xgboost.py
@@ -2,17 +2,29 @@
 import os
 import xgboost as xgb
 
+from typing import List
+
+from mlserver.errors import InferenceError
 from mlserver.settings import ModelSettings
 from mlserver.codecs import CodecError
-from mlserver.types import RequestInput, InferenceRequest
+from mlserver.types import RequestOutput, RequestInput, InferenceRequest
 
 from mlserver_xgboost import XGBoostModel
-from mlserver_xgboost.xgboost import WELLKNOWN_MODEL_FILENAMES
+from mlserver_xgboost.xgboost import (
+    WELLKNOWN_MODEL_FILENAMES,
+    PREDICT_OUTPUT,
+    PREDICT_PROBA_OUTPUT,
+)
 
 
 def test_load(model: XGBoostModel):
     assert model.ready
-    assert type(model._model) == xgb.Booster
+    assert type(model._model) == xgb.XGBRegressor
+
+
+def test_load_classifier(classifier: XGBoostModel):
+    assert classifier.ready
+    assert type(classifier._model) == xgb.XGBClassifier
 
 
 @pytest.mark.parametrize("fname", WELLKNOWN_MODEL_FILENAMES)
@@ -27,7 +39,7 @@ async def test_load_folder(fname, model_uri: str, model_settings: ModelSettings)
     await model.load()
 
     assert model.ready
-    assert type(model._model) == xgb.Booster
+    assert type(model._model) == xgb.XGBRegressor
 
 
 async def test_predict(model: XGBoostModel, inference_request: InferenceRequest):
@@ -37,6 +49,46 @@ async def test_predict(model: XGBoostModel, inference_request: InferenceRequest)
     assert 0 <= response.outputs[0].data[0] <= 1
 
 
+@pytest.mark.parametrize(
+    "req_outputs",
+    [
+        [],
+        [PREDICT_OUTPUT],
+        [PREDICT_PROBA_OUTPUT],
+        [PREDICT_OUTPUT, PREDICT_PROBA_OUTPUT],
+    ],
+)
+async def test_predict_classifier(
+    classifier: XGBoostModel,
+    inference_request: InferenceRequest,
+    req_outputs: List[str],
+):
+    inference_request.outputs = [
+        RequestOutput(name=req_output) for req_output in req_outputs
+    ]
+    response = await classifier.predict(inference_request)
+
+    input_data = inference_request.inputs[0].data
+    if len(req_outputs) == 0:
+        # Assert that PREDICT_OUTPUT is added by default
+        req_outputs = [PREDICT_OUTPUT]
+
+    assert len(response.outputs) == len(req_outputs)
+    for req_output, output in zip(req_outputs, response.outputs):
+        assert output.name == req_output
+        assert output.shape[0] == len(input_data)  # type: ignore
+
+
+@pytest.mark.parametrize("invalid_output", ["something_else", PREDICT_PROBA_OUTPUT])
+async def test_invalid_output_error(
+    model: XGBoostModel, inference_request: InferenceRequest, invalid_output: str
+):
+    inference_request.outputs = [RequestOutput(name=invalid_output)]
+
+    with pytest.raises(InferenceError):
+        await model.predict(inference_request)
+
+
 async def test_multiple_inputs_error(
     model: XGBoostModel, inference_request: InferenceRequest
 ):