Repository: kserve/kserve

Issue:
Support .json and .ubj model format for XGBoost server image
/kind feature


**Description**
In the XGBoost image, the only supported model format is .bst: https://github.com/kserve/kserve/blob/56b8fe0d189fc0d557e9a8af07eab0c12852d5fd/python/xgbserver/xgbserver/model.py#L28

This format has been deprecated for a while and is not backwards compatible between xgboost framework versions. The recommended model format is .json or .ubj: https://xgboost.readthedocs.io/en/stable/tutorials/saving_model.html

Users that want to use the recommended model format for XGBoost models, are currently not able to do so.


**Proposed solution**
Support the recommended file formats, while also keeping support for the old .bst format. 







Discussion:

Resolving patch:
diff --git a/python/xgbserver/xgbserver/model.py b/python/xgbserver/xgbserver/model.py
--- a/python/xgbserver/xgbserver/model.py
+++ b/python/xgbserver/xgbserver/model.py
@@ -25,7 +25,7 @@
 from kserve import Model
 from kserve.storage import Storage
 
-BOOSTER_FILE_EXTENSION = ".bst"
+BOOSTER_FILE_EXTENSIONS = (".bst", ".json", ".ubj")
 
 
 class XGBoostModel(Model):
@@ -45,7 +45,7 @@ def load(self) -> bool:
         model_files = []
         for file in os.listdir(model_path):
             file_path = os.path.join(model_path, file)
-            if os.path.isfile(file_path) and file.endswith(BOOSTER_FILE_EXTENSION):
+            if os.path.isfile(file_path) and file.endswith(BOOSTER_FILE_EXTENSIONS):
                 model_files.append(file_path)
         if len(model_files) == 0:
             raise ModelMissingError(model_path)


Test patch:
diff --git a/python/xgbserver/xgbserver/test_model.py b/python/xgbserver/xgbserver/test_model.py
--- a/python/xgbserver/xgbserver/test_model.py
+++ b/python/xgbserver/xgbserver/test_model.py
@@ -18,12 +18,16 @@
 from sklearn.datasets import load_iris
 from xgbserver import XGBoostModel
 
-model_dir = os.path.join(os.path.dirname(__file__), "example_model", "model")
+bst_model_dir = os.path.join(os.path.dirname(__file__), "example_model", "bst_model")
+json_model_dir = os.path.join(os.path.dirname(__file__), "example_model", "json_model")
+ubj_model_dir = os.path.join(os.path.dirname(__file__), "example_model", "ubj_model")
 BST_FILE = "model.bst"
+JSON_FILE = "model.json"
+UBJ_FILE = "model.ubj"
 NTHREAD = 1
 
 
-def test_model():
+def test_bst_model():
     iris = load_iris()
     y = iris["target"]
     X = iris["data"]
@@ -37,9 +41,71 @@ def test_model():
         "objective": "multi:softmax",
     }
     xgb_model = xgb.train(params=param, dtrain=dtrain)
-    model_file = os.path.join(model_dir, BST_FILE)
+    model_file = os.path.join(bst_model_dir, BST_FILE)
     xgb_model.save_model(model_file)
-    model = XGBoostModel("model", model_dir, NTHREAD)
+    model = XGBoostModel("model", bst_model_dir, NTHREAD)
+    model.load()
+    request = [X[0].tolist()]
+    response = model.predict({"instances": request})
+    assert response["predictions"] == [0]
+
+    # test v2 infer call
+    infer_input = InferInput(
+        name="input-0", shape=[1, 4], datatype="FP32", data=request
+    )
+    infer_request = InferRequest(model_name="model", infer_inputs=[infer_input])
+    infer_response = model.predict(infer_request)
+    assert infer_response.to_rest()["outputs"][0]["data"] == [0]
+
+
+def test_json_model():
+    iris = load_iris()
+    y = iris["target"]
+    X = iris["data"]
+    dtrain = xgb.DMatrix(X, label=y)
+    param = {
+        "max_depth": 6,
+        "eta": 0.1,
+        "silent": 1,
+        "nthread": 4,
+        "num_class": 10,
+        "objective": "multi:softmax",
+    }
+    xgb_model = xgb.train(params=param, dtrain=dtrain)
+    model_file = os.path.join(json_model_dir, JSON_FILE)
+    xgb_model.save_model(model_file)
+    model = XGBoostModel("model", json_model_dir, NTHREAD)
+    model.load()
+    request = [X[0].tolist()]
+    response = model.predict({"instances": request})
+    assert response["predictions"] == [0]
+
+    # test v2 infer call
+    infer_input = InferInput(
+        name="input-0", shape=[1, 4], datatype="FP32", data=request
+    )
+    infer_request = InferRequest(model_name="model", infer_inputs=[infer_input])
+    infer_response = model.predict(infer_request)
+    assert infer_response.to_rest()["outputs"][0]["data"] == [0]
+
+
+def test_ubj_model():
+    iris = load_iris()
+    y = iris["target"]
+    X = iris["data"]
+    dtrain = xgb.DMatrix(X, label=y)
+    param = {
+        "max_depth": 6,
+        "eta": 0.1,
+        "silent": 1,
+        "nthread": 4,
+        "num_class": 10,
+        "objective": "multi:softmax",
+    }
+    xgb_model = xgb.train(params=param, dtrain=dtrain)
+    model_file = os.path.join(ubj_model_dir, UBJ_FILE)
+    xgb_model.save_model(model_file)
+    model = XGBoostModel("model", ubj_model_dir, NTHREAD)
     model.load()
     request = [X[0].tolist()]
     response = model.predict({"instances": request})
diff --git a/python/xgbserver/xgbserver/test_xgboost_model_repository.py b/python/xgbserver/xgbserver/test_xgboost_model_repository.py
--- a/python/xgbserver/xgbserver/test_xgboost_model_repository.py
+++ b/python/xgbserver/xgbserver/test_xgboost_model_repository.py
@@ -17,13 +17,15 @@
 from xgbserver import XGBoostModelRepository
 
 model_dir = os.path.join(os.path.dirname(__file__), "example_model")
-invalid_model_dir = os.path.join(os.path.dirname(__file__), "model_not_exist", "model")
+invalid_model_dir = os.path.join(
+    os.path.dirname(__file__), "model_not_exist", "bst_model"
+)
 
 
 @pytest.mark.asyncio
 async def test_load():
     repo = XGBoostModelRepository(model_dir=model_dir, nthread=1)
-    model_name = "model"
+    model_name = "bst_model"
     await repo.load(model_name)
     assert repo.get_model(model_name) is not None
     assert repo.is_model_ready(model_name)
@@ -33,6 +35,6 @@ async def test_load():
 async def test_load_fail():
     with pytest.raises(Exception):
         repo = XGBoostModelRepository(model_dir=invalid_model_dir, nthread=1)
-        model_name = "model"
+        model_name = "bst_model"
         assert repo.get_model(model_name) is None
         assert not repo.is_model_ready(model_name)