[dask] Ensure model can be pickled. (#6651)

2021-01-28 21:47:57 +08:00
parent 0ad6e18a2a
commit d167892c7e
2 changed files with 20 additions and 5 deletions
--- 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.'''