Add tests for pickling with custom obj and metric. (#9943)

tests/python/test_pickling.py

@@ -1,10 +1,13 @@
 import json
 import os
 import pickle
+import tempfile
 
 import numpy as np
+import pytest
 
 import xgboost as xgb
+from xgboost import testing as tm
 
 kRows = 100
 kCols = 10
@@ -61,3 +64,27 @@ class TestPickling:
         params = {"nthread": 8, "tree_method": "exact", "subsample": 0.5}
         config = self.run_model_pickling(params)
         check(config)
+
+    @pytest.mark.skipif(**tm.no_sklearn())
+    def test_with_sklearn_obj_metric(self) -> None:
+        from sklearn.metrics import mean_squared_error
+
+        X, y = tm.datasets.make_regression()
+        reg = xgb.XGBRegressor(objective=tm.ls_obj, eval_metric=mean_squared_error)
+        reg.fit(X, y)
+
+        pkl = pickle.dumps(reg)
+        reg_1 = pickle.loads(pkl)
+        assert callable(reg_1.objective)
+        assert callable(reg_1.eval_metric)
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            path = os.path.join(tmpdir, "model.json")
+            reg.save_model(path)
+
+            reg_2 = xgb.XGBRegressor()
+            reg_2.load_model(path)
+
+        assert not callable(reg_2.objective)
+        assert not callable(reg_2.eval_metric)
+        assert reg_2.eval_metric is None

tests/python/test_with_sklearn.py

@@ -504,15 +504,10 @@ def test_regression_with_custom_objective():
     from sklearn.metrics import mean_squared_error
     from sklearn.model_selection import KFold
 
-    def objective_ls(y_true, y_pred):
-        grad = (y_pred - y_true)
-        hess = np.ones(len(y_true))
-        return grad, hess
-
     X, y = fetch_california_housing(return_X_y=True)
     kf = KFold(n_splits=2, shuffle=True, random_state=rng)
     for train_index, test_index in kf.split(X, y):
-        xgb_model = xgb.XGBRegressor(objective=objective_ls).fit(
+        xgb_model = xgb.XGBRegressor(objective=tm.ls_obj).fit(
             X[train_index], y[train_index]
         )
         preds = xgb_model.predict(X[test_index])
@@ -530,27 +525,29 @@ def test_regression_with_custom_objective():
     np.testing.assert_raises(XGBCustomObjectiveException, xgb_model.fit, X, y)
 
 
+def logregobj(y_true, y_pred):
+    y_pred = 1.0 / (1.0 + np.exp(-y_pred))
+    grad = y_pred - y_true
+    hess = y_pred * (1.0 - y_pred)
+    return grad, hess
+
+
 def test_classification_with_custom_objective():
     from sklearn.datasets import load_digits
     from sklearn.model_selection import KFold
 
-    def logregobj(y_true, y_pred):
-        y_pred = 1.0 / (1.0 + np.exp(-y_pred))
-        grad = y_pred - y_true
-        hess = y_pred * (1.0 - y_pred)
-        return grad, hess
-
     digits = load_digits(n_class=2)
-    y = digits['target']
-    X = digits['data']
+    y = digits["target"]
+    X = digits["data"]
     kf = KFold(n_splits=2, shuffle=True, random_state=rng)
     for train_index, test_index in kf.split(X, y):
         xgb_model = xgb.XGBClassifier(objective=logregobj)
         xgb_model.fit(X[train_index], y[train_index])
         preds = xgb_model.predict(X[test_index])
         labels = y[test_index]
-        err = sum(1 for i in range(len(preds))
-                  if int(preds[i] > 0.5) != labels[i]) / float(len(preds))
+        err = sum(
+            1 for i in range(len(preds)) if int(preds[i] > 0.5) != labels[i]
+        ) / float(len(preds))
         assert err < 0.1
 
     # Test that the custom objective function is actually used