Repository: mozilla/bugbug

Issue:
Use `Booster.save_model` when saving XGBoost models instead of `pickle`
```
WARNING: /Users/runner/work/xgboost/xgboost/python-package/build/temp.macosx-11.0-arm64-cpython-38/xgboost/src/learner.cc:553: 
  If you are loading a serialized model (like pickle in Python, RDS in R) generated by
  older XGBoost, please export the model by calling `Booster.save_model` from that version
  first, then load it back in current version. See:

    https://xgboost.readthedocs.io/en/latest/tutorials/saving_model.html

  for more details about differences between saving model and serializing.
```


Discussion:
Hi @suhaibmujahid, I'm an Outreachy Intern eager to contribute to this project by addressing some of the identified issues. To get started, I'm in the process of familiarizing myself with the project's structure. Would you kindly assist me by indicating which script contains the code for the XGBoost model? Your guidance would be greatly appreciated.
@Inyrkz thank you for your interest in the project! Please have a look at the readme file to get an idea about the project structure: https://github.com/mozilla/bugbug/tree/master#structure-of-the-project

@suhaibmujahid Thank you for steering me in the right direction.
Hello @suhaibmujahid I'm trying to work on this issue since no pr is linked to it! What prompted the above error?
The warning above appears when loading a model that was trained in an older version of XGBoost. For example, training a model using XGBoost v1.7.5, saving the model using `pickle`, then loading the model again after updating XGBoost to v1.7.6.



Resolving patch:
diff --git a/bugbug/model.py b/bugbug/model.py
--- a/bugbug/model.py
+++ b/bugbug/model.py
@@ -6,6 +6,7 @@
 import logging
 import pickle
 from collections import defaultdict
+from os import makedirs, path
 from typing import Any
 
 import matplotlib
@@ -24,6 +25,7 @@
 from sklearn.model_selection import cross_validate, train_test_split
 from sklearn.preprocessing import LabelEncoder
 from tabulate import tabulate
+from xgboost import XGBModel
 
 from bugbug import bugzilla, db, repository
 from bugbug.github import Github
@@ -567,7 +569,19 @@ def train(self, importance_cutoff=0.15, limit=None):
 
             self.clf.fit(X_train, self.le.transform(y_train))
 
-        with open(self.__class__.__name__.lower(), "wb") as f:
+        model_directory = self.__class__.__name__.lower()
+        makedirs(model_directory, exist_ok=True)
+
+        if issubclass(type(self.clf), XGBModel):
+            xgboost_model_path = path.join(model_directory, "xgboost.ubj")
+            self.clf.save_model(xgboost_model_path)
+
+            # Since we save the classifier separately, we need to clear the clf
+            # attribute to prevent it from being pickled with the model object.
+            self.clf = self.clf.__class__(n_jobs=self.clf.n_jobs)
+
+        model_path = path.join(model_directory, "model.pkl")
+        with open(model_path, "wb") as f:
             pickle.dump(self, f, protocol=pickle.HIGHEST_PROTOCOL)
 
         if self.store_dataset:
@@ -580,9 +594,16 @@ def train(self, importance_cutoff=0.15, limit=None):
         return tracking_metrics
 
     @staticmethod
-    def load(model_file_name: str) -> "Model":
-        with open(model_file_name, "rb") as f:
-            return pickle.load(f)
+    def load(model_directory: str) -> "Model":
+        model_path = path.join(model_directory, "model.pkl")
+        with open(model_path, "rb") as f:
+            model = pickle.load(f)
+
+        xgboost_model_path = path.join(model_directory, "xgboost.ubj")
+        if path.exists(xgboost_model_path):
+            model.clf.load_model(xgboost_model_path)
+
+        return model
 
     def overwrite_classes(self, items, classes, probabilities):
         return classes
diff --git a/bugbug/utils.py b/bugbug/utils.py
--- a/bugbug/utils.py
+++ b/bugbug/utils.py
@@ -231,13 +231,13 @@ def download_model(model_name: str) -> str:
             version = "latest"
 
     path = f"{model_name}model"
-    url = f"https://community-tc.services.mozilla.com/api/index/v1/task/project.bugbug.train_{model_name}.{version}/artifacts/public/{path}.zst"
+    url = f"https://community-tc.services.mozilla.com/api/index/v1/task/project.bugbug.train_{model_name}.{version}/artifacts/public/{path}.tar.zst"
     logger.info("Downloading %s...", url)
     updated = download_check_etag(url)
     if updated:
-        zstd_decompress(path)
-        os.remove(f"{path}.zst")
-    assert os.path.exists(path), "Decompressed file exists"
+        extract_tar_zst(f"{path}.tar.zst")
+        os.remove(f"{path}.tar.zst")
+    assert os.path.exists(path), "Decompressed directory exists"
     return path
 
 
diff --git a/scripts/trainer.py b/scripts/trainer.py
--- a/scripts/trainer.py
+++ b/scripts/trainer.py
@@ -9,7 +9,7 @@
 
 from bugbug import db
 from bugbug.models import MODELS, get_model_class
-from bugbug.utils import CustomJsonEncoder, zstd_compress
+from bugbug.utils import CustomJsonEncoder, create_tar_zst, zstd_compress
 
 MODELS_WITH_TYPE = ("component",)
 
@@ -57,17 +57,17 @@ def go(self, args):
 
         logger.info("Training done")
 
-        model_file_name = f"{model_name}model"
-        assert os.path.exists(model_file_name)
-        zstd_compress(model_file_name)
+        model_directory = f"{model_name}model"
+        assert os.path.exists(model_directory)
+        create_tar_zst(f"{model_directory}.tar.zst")
 
         logger.info("Model compressed")
 
         if model_obj.store_dataset:
-            assert os.path.exists(f"{model_file_name}_data_X")
-            zstd_compress(f"{model_file_name}_data_X")
-            assert os.path.exists(f"{model_file_name}_data_y")
-            zstd_compress(f"{model_file_name}_data_y")
+            assert os.path.exists(f"{model_name}model_data_X")
+            zstd_compress(f"{model_name}model_data_X")
+            assert os.path.exists(f"{model_name}model_data_y")
+            zstd_compress(f"{model_name}model_data_y")
 
 
 def parse_args(args):


Test patch: