diff --git a/python-package/xgboost/dask.py b/python-package/xgboost/dask.py index 2cae5fc5c..1927ba54a 100644 --- a/python-package/xgboost/dask.py +++ b/python-package/xgboost/dask.py @@ -1359,6 +1359,12 @@ class DaskScikitLearnBase(XGBModel): return self return self.client.sync(_).__await__() + def __getstate__(self): + this = self.__dict__.copy() + if "_client" in this.keys(): + del this["_client"] + return this + @property def client(self) -> "distributed.Client": '''The dask client used in this model.''' diff --git a/tests/python/test_with_dask.py b/tests/python/test_with_dask.py index 4861e19da..d58404960 100644 --- a/tests/python/test_with_dask.py +++ b/tests/python/test_with_dask.py @@ -1,5 +1,5 @@ from pathlib import Path - +import pickle import testing as tm import pytest import xgboost as xgb @@ -1104,23 +1104,32 @@ class TestWithDask: predt_0 = cls.predict(X) with tempfile.TemporaryDirectory() as tmpdir: + path = os.path.join(tmpdir, "model.pkl") + with open(path, "wb") as fd: + pickle.dump(cls, fd) + + with open(path, "rb") as fd: + cls = pickle.load(fd) + predt_1 = cls.predict(X) + np.testing.assert_allclose(predt_0.compute(), predt_1.compute()) + path = os.path.join(tmpdir, 'cls.json') cls.save_model(path) cls = xgb.dask.DaskXGBClassifier() cls.load_model(path) assert cls.n_classes_ == 10 - predt_1 = cls.predict(X) + predt_2 = cls.predict(X) - np.testing.assert_allclose(predt_0.compute(), predt_1.compute()) + np.testing.assert_allclose(predt_0.compute(), predt_2.compute()) # Use single node to load cls = xgb.XGBClassifier() cls.load_model(path) assert cls.n_classes_ == 10 - predt_2 = cls.predict(X_) + predt_3 = cls.predict(X_) - np.testing.assert_allclose(predt_0.compute(), predt_2) + np.testing.assert_allclose(predt_0.compute(), predt_3) class TestDaskCallbacks: