diff --git a/python/cuml/cuml/dask/common/base.py b/python/cuml/cuml/dask/common/base.py index 1f2f71542c..96806a6191 100644 --- a/python/cuml/cuml/dask/common/base.py +++ b/python/cuml/cuml/dask/common/base.py @@ -57,7 +57,9 @@ def __init__(self, *, client=None, verbose=False, **kwargs): self.internal_model = None def __getstate__(self): - internal_model = self._get_internal_model().result() + internal_model = self._get_internal_model() + if internal_model: + internal_model = internal_model.result() state = { "verbose": self.verbose, "kwargs": self.kwargs, diff --git a/python/cuml/cuml/tests/dask/test_dask_serialization.py b/python/cuml/cuml/tests/dask/test_dask_serialization.py index 3dc819b24f..0abc822be5 100644 --- a/python/cuml/cuml/tests/dask/test_dask_serialization.py +++ b/python/cuml/cuml/tests/dask/test_dask_serialization.py @@ -80,3 +80,16 @@ def test_serialize_mnmg_model(client): unpickled_model = pickle.loads(pickled_model) assert np.allclose(unpickled_model.coef_, model.coef_) + + +def test_serialize_before_training(client): + X, y = make_regression(n_samples=1000, n_features=20, random_state=0) + X, y = da.from_array(X), da.from_array(y) + + model = LinearRegression(client=client) + pickled_model = pickle.dumps(model) + unpickled_model = pickle.loads(pickled_model) + + unpickled_model.client = client + unpickled_model.fit(X, y) + assert hasattr(unpickled_model, "coef_")