Update predict_utils.py

predict_utils.py

@@ -1,5 +1,5 @@
 # predict_utils.py
-# Robust loader + monkey-patches for XGBoost sklearn wrappers
+# Robust loader + extended monkey-patches for XGBoost sklearn wrappers
 import os
 import logging
 import joblib
@@ -21,11 +21,9 @@ LOCAL_CANDIDATES = [
 ]
 
 # -------------------------
-# Monkey-patch xgboost sklearn wrappers & base class to add missing attributes.
-# This handles errors like:
-#   "'XGBClassifier' object has no attribute 'use_label_encoder'"
-#   "'XGBModel' object has no attribute 'gpu_id'"
-# Call this BEFORE joblib.load so unpickling has the attributes available.
+# Extended monkey-patch
+# Add commonly-expected attributes so unpickling older models succeeds.
+# Call this BEFORE joblib.load so unpickle finds these attributes.
 # -------------------------
 def ensure_xgb_sklearn_compat():
     try:
@@ -34,66 +32,54 @@ def ensure_xgb_sklearn_compat():
         logger.debug(f"xgboost not importable for patching: {e}")
         return
 
-    # Base class: XGBModel (add gpu_id if missing)
+    # Attributes to add on XGBModel base class (safe defaults)
     XGBModel = getattr(xgb, "XGBModel", None)
     if XGBModel is not None:
-        if not hasattr(XGBModel, "gpu_id"):
+        for attr, val in {
+            "gpu_id": None,
+            "nthread": None,
+            "n_jobs": None,
+            "predictor": None,
+            "base_score": None,
+            "objective": None,
+        }.items():
             try:
-                setattr(XGBModel, "gpu_id", None)  # default None means CPU
-                logger.info("Patched XGBModel.gpu_id = None")
+                if not hasattr(XGBModel, attr):
+                    setattr(XGBModel, attr, val)
+                    logger.info(f"Patched XGBModel.{attr} = {val!r}")
             except Exception as e:
-                logger.debug(f"Could not set XGBModel.gpu_id: {e}")
-
-    # XGBClassifier: add use_label_encoder if missing, maybe objective
-    XGBClassifier = getattr(xgb, "XGBClassifier", None)
-    if XGBClassifier is not None:
-        # Add attributes on the class object so unpickle will find defaults
-        if not hasattr(XGBClassifier, "use_label_encoder"):
-            try:
-                setattr(XGBClassifier, "use_label_encoder", False)
-                logger.info("Patched XGBClassifier.use_label_encoder = False")
-            except Exception as e:
-                logger.debug(f"Could not patch XGBClassifier.use_label_encoder: {e}")
-        # sometimes older pickles expect 'objective' attribute default
-        if not hasattr(XGBClassifier, "objective"):
-            try:
-                setattr(XGBClassifier, "objective", None)
-                logger.info("Patched XGBClassifier.objective = None")
-            except Exception as e:
-                logger.debug(f"Could not patch XGBClassifier.objective: {e}")
-
-    # XGBRegressor: similar patches
-    XGBRegressor = getattr(xgb, "XGBRegressor", None)
-    if XGBRegressor is not None:
-        if not hasattr(XGBRegressor, "use_label_encoder"):
-            try:
-                setattr(XGBRegressor, "use_label_encoder", False)
-                logger.info("Patched XGBRegressor.use_label_encoder = False")
-            except Exception as e:
-                logger.debug(f"Could not patch XGBRegressor.use_label_encoder: {e}")
-        if not hasattr(XGBRegressor, "objective"):
-            try:
-                setattr(XGBRegressor, "objective", None)
-                logger.info("Patched XGBRegressor.objective = None")
-            except Exception as e:
-                logger.debug(f"Could not patch XGBRegressor.objective: {e}")
-
-    # Also handle the case where pickled objects expect 'nthread' or 'n_jobs'
-    if XGBModel is not None:
-        if not hasattr(XGBModel, "nthread"):
-            try:
-                setattr(XGBModel, "nthread", None)
-                logger.info("Patched XGBModel.nthread = None")
-            except Exception as e:
-                logger.debug(f"Could not patch XGBModel.nthread: {e}")
-        if not hasattr(XGBModel, "n_jobs"):
-            try:
-                setattr(XGBModel, "n_jobs", None)
-                logger.info("Patched XGBModel.n_jobs = None")
-            except Exception as e:
-                logger.debug(f"Could not patch XGBModel.n_jobs: {e}")
+                logger.debug(f"Could not patch XGBModel.{attr}: {e}")
+
+    # Patch classifier/regressor class-level defaults used in older pickles
+    for cls_name in ("XGBClassifier", "XGBRegressor"):
+        cls = getattr(xgb, cls_name, None)
+        if cls is not None:
+            for attr, val in {
+                "use_label_encoder": False,
+                "objective": None,
+                "predictor": None,
+                "gpu_id": None,
+                "n_jobs": None,
+                "nthread": None,
+            }.items():
+                try:
+                    if not hasattr(cls, attr):
+                        setattr(cls, attr, val)
+                        logger.info(f"Patched {cls_name}.{attr} = {val!r}")
+                except Exception as e:
+                    logger.debug(f"Could not patch {cls_name}.{attr}: {e}")
 
-# Call patch early so joblib.load can use these defaults
+    # Some pickles expect certain module-level names (rare) — leave safe no-op fallbacks
+    try:
+        # e.g., older pickles might refer to xgb.core.Booster attributes; skip if not present
+        core = getattr(xgb, "core", None)
+        if core is not None:
+            if not hasattr(core, "DataBatch"):
+                setattr(core, "DataBatch", object)
+    except Exception:
+        pass
+
+# Run the patch early
 ensure_xgb_sklearn_compat()
 
 # -------------------------
@@ -119,7 +105,7 @@ def inspect_file(path):
 
 def try_joblib_load(path):
     try:
-        # Ensure patch right before load (in case xgboost gets imported lazily)
+        # ensure patch immediately before load (handles lazy imports)
         ensure_xgb_sklearn_compat()
         logger.info(f"Trying joblib.load on {path}")
         m = joblib.load(path)