[breaking] Change internal model serialization to UBJSON. (#7556)

* Use typed array for models.
* Change the memory snapshot format.
* Add new C API for saving to raw format.

tests/python/test_basic_models.py

@@ -14,7 +14,7 @@ dtest = xgb.DMatrix(dpath + 'agaricus.txt.test')
 rng = np.random.RandomState(1994)
 
 
-def json_model(model_path, parameters):
+def json_model(model_path: str, parameters: dict) -> dict:
     X = np.random.random((10, 3))
     y = np.random.randint(2, size=(10,))
 
@@ -22,9 +22,14 @@ def json_model(model_path, parameters):
 
     bst = xgb.train(parameters, dm1)
     bst.save_model(model_path)
+    if model_path.endswith("ubj"):
+        import ubjson
+        with open(model_path, "rb") as ubjfd:
+            model = ubjson.load(ubjfd)
+    else:
+        with open(model_path, 'r') as fd:
+            model = json.load(fd)
 
-    with open(model_path, 'r') as fd:
-        model = json.load(fd)
     return model
 
 
@@ -259,23 +264,40 @@ class TestModels:
         buf_from_raw = from_raw.save_raw()
         assert buf == buf_from_raw
 
-    def test_model_json_io(self):
+    def run_model_json_io(self, parameters: dict, ext: str) -> None:
+        if ext == "ubj" and tm.no_ubjson()["condition"]:
+            pytest.skip(tm.no_ubjson()["reason"])
+
         loc = locale.getpreferredencoding(False)
-        model_path = 'test_model_json_io.json'
-        parameters = {'tree_method': 'hist', 'booster': 'gbtree'}
+        model_path = 'test_model_json_io.' + ext
         j_model = json_model(model_path, parameters)
         assert isinstance(j_model['learner'], dict)
 
         bst = xgb.Booster(model_file=model_path)
 
         bst.save_model(fname=model_path)
-        with open(model_path, 'r') as fd:
-            j_model = json.load(fd)
+        if ext == "ubj":
+            import ubjson
+            with open(model_path, "rb") as ubjfd:
+                j_model = ubjson.load(ubjfd)
+        else:
+            with open(model_path, 'r') as fd:
+                j_model = json.load(fd)
+
         assert isinstance(j_model['learner'], dict)
 
         os.remove(model_path)
         assert locale.getpreferredencoding(False) == loc
 
+    @pytest.mark.parametrize("ext", ["json", "ubj"])
+    def test_model_json_io(self, ext: str) -> None:
+        parameters = {"booster": "gbtree", "tree_method": "hist"}
+        self.run_model_json_io(parameters, ext)
+        parameters = {"booster": "gblinear"}
+        self.run_model_json_io(parameters, ext)
+        parameters = {"booster": "dart", "tree_method": "hist"}
+        self.run_model_json_io(parameters, ext)
+
     @pytest.mark.skipif(**tm.no_json_schema())
     def test_json_io_schema(self):
         import jsonschema

tests/python/test_pickling.py

@@ -2,6 +2,7 @@ import pickle
 import numpy as np
 import xgboost as xgb
 import os
+import json
 
 
 kRows = 100
@@ -15,13 +16,14 @@ def generate_data():
 
 
 class TestPickling:
-    def run_model_pickling(self, xgb_params):
+    def run_model_pickling(self, xgb_params) -> str:
         X, y = generate_data()
         dtrain = xgb.DMatrix(X, y)
         bst = xgb.train(xgb_params, dtrain)
 
         dump_0 = bst.get_dump(dump_format='json')
         assert dump_0
+        config_0 = bst.save_config()
 
         filename = 'model.pkl'
 
@@ -42,9 +44,22 @@ class TestPickling:
         if os.path.exists(filename):
             os.remove(filename)
 
+        config_1 = bst.save_config()
+        assert config_0 == config_1
+        return json.loads(config_0)
+
     def test_model_pickling_json(self):
-        params = {
-            'nthread': 1,
-            'tree_method': 'hist',
-        }
-        self.run_model_pickling(params)
+        def check(config):
+            updater = config["learner"]["gradient_booster"]["updater"]
+            if params["tree_method"] == "exact":
+                subsample = updater["grow_colmaker"]["train_param"]["subsample"]
+            else:
+                subsample = updater["grow_quantile_histmaker"]["train_param"]["subsample"]
+            assert float(subsample) == 0.5
+
+        params = {"nthread": 8, "tree_method": "hist", "subsample": 0.5}
+        config = self.run_model_pickling(params)
+        check(config)
+        params = {"nthread": 8, "tree_method": "exact", "subsample": 0.5}
+        config = self.run_model_pickling(params)
+        check(config)

tests/python/testing.py

@@ -29,6 +29,15 @@ except ImportError:
 memory = Memory('./cachedir', verbose=0)
 
 
+def no_ubjson():
+    reason = "ubjson is not intsalled."
+    try:
+        import ubjson           # noqa
+        return {"condition": False, "reason": reason}
+    except ImportError:
+        return {"condition": True, "reason": reason}
+
+
 def no_sklearn():
     return {'condition': not SKLEARN_INSTALLED,
             'reason': 'Scikit-Learn is not installed'}