Repository: SeldonIO/MLServer

Issue:
Expected XGBoost model file "model.bst" extension is undocumented? 
On https://github.com/SeldonIO/MLServer/blob/master/runtimes/xgboost/mlserver_xgboost/xgboost.py#L21 you can see that MLServer is looking for an XGBoost model file called "model.bst". However, I cannot find any reference to that file extension in the XGBoost documentation. As far as I can see, XGBoost's documented file extensions are:

- ".json" added in 1.0.0, an "open format that can be easily reused"
- ".ubj" for Universal Binary JSON format, available in 1.6.0
- ".model" for the "old binary internal format" prior to 1.0.0, as shown in examples

Where does MLServer get the ".bst" extension from, and what model format does it use? Shouldn't it use one of the extensions mentioned in the XGBoost documentation instead, to avoid ambiguity?


Discussion:
Hey @rehevkor5 ,

That's a great point. 

We mainly support the `*.bst` extension to keep compatibility with the older "non-v2 protocol" XGBoost servers in both Seldon Core and KServe. But TBH I don't know where the name comes from. Having said that, it seems that other solutions like [Vertex AI](https://cloud.google.com/vertex-ai/docs/training/exporting-model-artifacts) also refer to this `*.bst` extension, so it may be a remnant from the time when XGBoost didn't support the JSON format?

Keep in mind though, that MLServer also supports the `*.json` extension. By default, if you pass a folder as the model's URI, it will try to look for a `model.bst` or `model.json` files. Alternatively, you can specify any file name directly on the `parameters.uri` field of the `model-settings.json` file. 

Haven't read much about the `.ubj` or `.model` extensions. Would it be worth also supporting those on the  `mlserver-xgboost` runtime? 
> TBH I don't know where the name comes from.
> it may be a remnant from the time when XGBoost didn't support the JSON format?

Gotcha. I tried looking at the XGBoost code.

I'm not a C++ expert (it's been a while), but https://github.com/dmlc/xgboost/blob/v1.5.0/src/c_api/c_api.cc#L843 and https://github.com/dmlc/xgboost/blob/v1.5.0/src/c_api/c_api.cc#L862 only seems to check for "json". Otherwise, I assume it's calling `virtual void SaveModel(dmlc::Stream* fo) const = 0;` from learner.h, which must be implemented here https://github.com/dmlc/xgboost/blob/v1.5.0/src/learner.cc#L876 So, it looks like any extension other than "json" will use the old binary format.

> Haven't read much about the .ubj or .model extensions. Would it be worth also supporting those on the mlserver-xgboost runtime?

Regarding `.ubj`, I cannot see in the documentation how to tell which version of XGBoost each version of MLServer supports/includes. However, `.ubj` is only available starting in XGBoost 1.6.0.

Is there some way to tell whether a specific MLServer version includes XGBoost 1.6.0? (Anything easier than, I suppose, opening an interactive shell into the Docker container & inspecting the libraries on disk?)
But yes, I agree that MLServer should support `.ubj` if possible.

Resolving patch:
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
@@ -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:


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
@@ -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
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
@@ -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()