Add support for .ubj models and improve XGBoost docs (SeldonIO#1168)

runtimes/xgboost/README.md

@@ -13,6 +13,37 @@ pip install mlserver mlserver-xgboost
 For further information on how to use MLServer with XGBoost, you can check out
 this [worked out example](../../docs/examples/xgboost/README.md).
 
+## XGBoost Artifact Type
+
+The XGBoost inference runtime will expect that your model is serialised via one
+of the following methods:
+
+| Extension | Docs                                                                                                                 | Example                            |
+| --------- | -------------------------------------------------------------------------------------------------------------------- | ---------------------------------- |
+| `*.json`  | [JSON Format](https://xgboost.readthedocs.io/en/stable/tutorials/saving_model.html#introduction-to-model-io)         | `booster.save_model("model.json")` |
+| `*.ubj`   | [Binary JSON Format](https://xgboost.readthedocs.io/en/stable/tutorials/saving_model.html#introduction-to-model-io)  | `booster.save_model("model.ubj")`  |
+| `*.bst`   | [(Old) Binary Format](https://xgboost.readthedocs.io/en/stable/tutorials/saving_model.html#introduction-to-model-io) | `booster.save_model("model.bst")`  |
+
+````{note}
+By default, the runtime will look for a file called `model.[json | ubj | bst]`.
+However, this can be modified through the `parameters.uri` field of your
+{class}`ModelSettings <mlserver.settings.ModelSettings>` config (see the
+section on [Model Settings](../../docs/reference/model-settings.md) for more
+details).
+
+```{code-block} json
+---
+emphasize-lines: 3-5
+---
+{
+  "name": "foo",
+  "parameters": {
+    "uri": "./my-own-model-filename.json"
+  }
+}
+```
+````
+
 ## Content Types
 
 If no [content type](../../docs/user-guide/content-type) is present on the
@@ -21,3 +52,42 @@ request or metadata, the XGBoost runtime will try to decode the payload as a
 To avoid this, either send a different content type explicitly, or define the
 correct one as part of your [model's
 metadata](../../docs/reference/model-settings).
+
+## Model Outputs
+
+The XGBoost inference runtime exposes a number of outputs depending on the
+model type.
+These outputs match to the `predict` and `predict_proba` methods of the XGBoost
+model.
+
+| Output          | Returned By Default | Availability                                                          |
+| --------------- | ------------------- | --------------------------------------------------------------------- |
+| `predict`       | ✅                  | Available on all XGBoost models.                                      |
+| `predict_proba` | ❌                  | Only available on non-regressor models (i.e. `XGBClassifier` models). |
+
+By default, the runtime will only return the output of `predict`.
+However, you are able to control which outputs you want back through the
+`outputs` field of your {class}`InferenceRequest
+<mlserver.types.InferenceRequest>` payload.
+
+For example, to only return the model's `predict_proba` output, you could
+define a payload such as:
+
+```{code-block} json
+---
+emphasize-lines: 10-12
+---
+{
+  "inputs": [
+    {
+      "name": "my-input",
+      "datatype": "INT32",
+      "shape": [2, 2],
+      "data": [1, 2, 3, 4]
+    }
+  ],
+  "outputs": [
+    { "name": "predict_proba" }
+  ]
+}
+```

runtimes/xgboost/mlserver_xgboost/xgboost.py

@@ -18,7 +18,7 @@
 PREDICT_PROBA_OUTPUT = "predict_proba"
 VALID_OUTPUTS = [PREDICT_OUTPUT, PREDICT_PROBA_OUTPUT]
 
-WELLKNOWN_MODEL_FILENAMES = ["model.bst", "model.json"]
+WELLKNOWN_MODEL_FILENAMES = ["model.bst", "model.json", "model.ubj"]
 
 
 def _load_sklearn_interface(model_uri: str) -> XGBModel:

runtimes/xgboost/tests/conftest.py

@@ -9,6 +9,7 @@
 from mlserver.utils import install_uvloop_event_loop
 
 from mlserver_xgboost import XGBoostModel
+from mlserver_xgboost.xgboost import WELLKNOWN_MODEL_FILENAMES
 
 TESTS_PATH = os.path.dirname(__file__)
 TESTDATA_PATH = os.path.join(TESTS_PATH, "testdata")
@@ -23,15 +24,16 @@ def event_loop():
     loop.close()
 
 
-@pytest.fixture
-def model_uri(tmp_path) -> str:
+@pytest.fixture(params=WELLKNOWN_MODEL_FILENAMES)
+def model_uri(request, tmp_path) -> str:
     n = 4
     d = 3
 
     dtrain = xgb.DMatrix(data=np.random.rand(n, d), label=np.random.rand(n))
     bst = xgb.train(params={}, dtrain=dtrain)
 
-    model_uri = os.path.join(tmp_path, "xgboost-model.json")
+    _, ext = os.path.splitext(request.param)
+    model_uri = os.path.join(tmp_path, f"xgboost-model{ext}")
     bst.save_model(model_uri)
 
     return model_uri

runtimes/xgboost/tests/test_xgboost.py

@@ -11,7 +11,6 @@
 
 from mlserver_xgboost import XGBoostModel
 from mlserver_xgboost.xgboost import (
-    WELLKNOWN_MODEL_FILENAMES,
     PREDICT_OUTPUT,
     PREDICT_PROBA_OUTPUT,
 )
@@ -27,13 +26,15 @@ def test_load_classifier(classifier: XGBoostModel):
     assert type(classifier._model) == xgb.XGBClassifier
 
 
-@pytest.mark.parametrize("fname", WELLKNOWN_MODEL_FILENAMES)
-async def test_load_folder(fname, model_uri: str, model_settings: ModelSettings):
+async def test_load_folder(model_uri: str, model_settings: ModelSettings):
+    # Rename `xgboost-model.[ext]` to `model.[ext]`
+    _, ext = os.path.splitext(model_uri)
+    fname = f"model{ext}"
     model_folder = os.path.dirname(model_uri)
     model_path = os.path.join(model_folder, fname)
     os.rename(model_uri, model_path)
 
-    model_settings.parameters.uri = model_path  # type: ignore
+    model_settings.parameters.uri = model_folder  # type: ignore
 
     model = XGBoostModel(model_settings)
     model.ready = await model.load()