Fix cutting edge CI failure (#276)

Fixes CI failure with cutting edge by ensuring the docstrings replaced in sklearn integration are searched for in superclasses.

Author: Yard1

xgboost_ray/sklearn.py

@@ -32,6 +32,7 @@
 import warnings
 import functools
 import inspect
+from inspect import _finddoc
 
 from ray.util.annotations import PublicAPI, DeveloperAPI
 
@@ -212,22 +213,41 @@ def _cls_predict_proba(n_classes: int, prediction, vstack: Callable):
 )
 
 
-def _treat_estimator_doc(doc: str) -> str:
+def _get_doc(object: Any) -> Optional[str]:
+    """Same as ``inspect.getdoc``, but without ``cleandoc`` applied."""
+    try:
+        doc = object.__doc__
+    except AttributeError:
+        return None
+    if doc is None:
+        try:
+            doc = _finddoc(object)
+        except (AttributeError, TypeError):
+            return None
+    if not isinstance(doc, str):
+        return None
+    return doc
+
+
+def _treat_estimator_doc(doc: Optional[str]) -> Optional[str]:
     """Helper function to make nececssary changes in estimator docstrings"""
-    doc = doc.replace(*_N_JOBS_DOC_REPLACE).replace(
-        "scikit-learn API for XGBoost",
-        "scikit-learn API for Ray-distributed XGBoost").replace(
-            ":doc:`tree method\n        </treemethod>`", "tree method")
+    if doc:
+        doc = doc.replace(*_N_JOBS_DOC_REPLACE).replace(
+            "scikit-learn API for XGBoost",
+            "scikit-learn API for Ray-distributed XGBoost").replace(
+                ":doc:`tree method\n        </treemethod>`", "tree method")
     return doc
 
 
-def _treat_X_doc(doc: str) -> str:
-    doc = doc.replace("Data to predict with.",
-                      "Data to predict with. Can also be a ``RayDMatrix``.")
-    doc = doc.replace("Feature matrix.",
-                      "Feature matrix. Can also be a ``RayDMatrix``.")
-    doc = doc.replace("Feature matrix",
-                      "Feature matrix. Can also be a ``RayDMatrix``.")
+def _treat_X_doc(doc: Optional[str]) -> Optional[str]:
+    if doc:
+        doc = doc.replace(
+            "Data to predict with.",
+            "Data to predict with. Can also be a ``RayDMatrix``.")
+        doc = doc.replace("Feature matrix.",
+                          "Feature matrix. Can also be a ``RayDMatrix``.")
+        doc = doc.replace("Feature matrix",
+                          "Feature matrix. Can also be a ``RayDMatrix``.")
     return doc
 
 
@@ -514,7 +534,7 @@ def fit(
         self._set_evaluation_result(evals_result)
         return self
 
-    fit.__doc__ = _treat_X_doc(XGBRegressor.fit.__doc__) + _RAY_PARAMS_DOC
+    fit.__doc__ = _treat_X_doc(_get_doc(XGBRegressor.fit)) + _RAY_PARAMS_DOC
 
     def _can_use_inplace_predict(self) -> bool:
         return False
@@ -542,16 +562,16 @@ def predict(
             _remote=_remote,
             ray_dmatrix_params=ray_dmatrix_params)
 
-    predict.__doc__ = _treat_X_doc(
-        XGBRegressor.predict.__doc__) + _RAY_PARAMS_DOC
+    predict.__doc__ = _treat_X_doc(_get_doc(
+        XGBRegressor.predict)) + _RAY_PARAMS_DOC
 
     def load_model(self, fname):
         if not hasattr(self, "_Booster"):
             self._Booster = Booster()
         return super().load_model(fname)
 
 
-RayXGBRegressor.__doc__ = _treat_estimator_doc(XGBRegressor.__doc__)
+RayXGBRegressor.__doc__ = _treat_estimator_doc(_get_doc(XGBRegressor))
 
 
 class RayXGBRFRegressor(RayXGBRegressor):
@@ -589,7 +609,7 @@ def get_num_boosting_rounds(self):
         return 1
 
 
-RayXGBRFRegressor.__doc__ = _treat_estimator_doc(XGBRFRegressor.__doc__)
+RayXGBRFRegressor.__doc__ = _treat_estimator_doc(_get_doc(XGBRFRegressor))
 
 
 @PublicAPI(stability="beta")
@@ -734,7 +754,7 @@ def fit(
         self._set_evaluation_result(evals_result)
         return self
 
-    fit.__doc__ = _treat_X_doc(XGBClassifier.fit.__doc__) + _RAY_PARAMS_DOC
+    fit.__doc__ = _treat_X_doc(_get_doc(XGBClassifier.fit)) + _RAY_PARAMS_DOC
 
     def _ray_fit_preprocess(self, y) -> Callable:
         """This has been separated out so that it can be easily overwritten
@@ -838,7 +858,8 @@ def predict(
             return self._le.inverse_transform(column_indexes)
         return column_indexes
 
-    predict.__doc__ = _treat_X_doc(XGBModel.predict.__doc__) + _RAY_PARAMS_DOC
+    predict.__doc__ = _treat_X_doc(_get_doc(
+        XGBModel.predict)) + _RAY_PARAMS_DOC
 
     def predict_proba(
             self,
@@ -872,10 +893,10 @@ def load_model(self, fname):
         return super().load_model(fname)
 
     predict_proba.__doc__ = (
-        _treat_X_doc(XGBClassifier.predict_proba.__doc__) + _RAY_PARAMS_DOC)
+        _treat_X_doc(_get_doc(XGBClassifier.predict_proba)) + _RAY_PARAMS_DOC)
 
 
-RayXGBClassifier.__doc__ = _treat_estimator_doc(XGBClassifier.__doc__)
+RayXGBClassifier.__doc__ = _treat_estimator_doc(_get_doc(XGBClassifier))
 
 
 class RayXGBRFClassifier(RayXGBClassifier):
@@ -935,7 +956,7 @@ def get_num_boosting_rounds(self):
         return 1
 
 
-RayXGBRFClassifier.__doc__ = _treat_estimator_doc(XGBRFClassifier.__doc__)
+RayXGBRFClassifier.__doc__ = _treat_estimator_doc(_get_doc(XGBRFClassifier))
 
 
 @PublicAPI(stability="beta")
@@ -1053,7 +1074,7 @@ def fit(
         self._set_evaluation_result(evals_result)
         return self
 
-    fit.__doc__ = _treat_X_doc(XGBRanker.fit.__doc__) + _RAY_PARAMS_DOC
+    fit.__doc__ = _treat_X_doc(_get_doc(XGBRanker.fit)) + _RAY_PARAMS_DOC
 
     def _can_use_inplace_predict(self) -> bool:
         return False
@@ -1081,12 +1102,13 @@ def predict(
             _remote=_remote,
             ray_dmatrix_params=ray_dmatrix_params)
 
-    predict.__doc__ = _treat_X_doc(XGBRanker.predict.__doc__) + _RAY_PARAMS_DOC
+    predict.__doc__ = _treat_X_doc(_get_doc(
+        XGBRanker.predict)) + _RAY_PARAMS_DOC
 
     def load_model(self, fname):
         if not hasattr(self, "_Booster"):
             self._Booster = Booster()
         return super().load_model(fname)
 
 
-RayXGBRanker.__doc__ = _treat_estimator_doc(XGBRanker.__doc__)
+RayXGBRanker.__doc__ = _treat_estimator_doc(_get_doc(XGBRanker))

xgboost_ray/tests/test_sklearn.py

@@ -976,7 +976,10 @@ def test_constraint_parameters(self):
         reg.fit(X, y)
 
         config = json.loads(reg.get_booster().save_config())
-        if XGBOOST_VERSION >= Version("1.6.0"):
+        if XGBOOST_VERSION > Version("1.7.4"):
+            assert (config["learner"]["gradient_booster"]["tree_train_param"][
+                "interaction_constraints"] == "[[0, 1], [2, 3, 4]]")
+        elif XGBOOST_VERSION >= Version("1.6.0"):
             assert (config["learner"]["gradient_booster"]["updater"][
                 "grow_histmaker"]["train_param"]["interaction_constraints"] ==
                     "[[0, 1], [2, 3, 4]]")