From 913eb97d43bf54bd21b75dd0d26ee11bd4d8e823 Mon Sep 17 00:00:00 2001 From: viclafargue Date: Tue, 10 Sep 2024 10:30:43 +0200 Subject: [PATCH 1/2] Fix Dask estimators serialization prior to training --- python/cuml/cuml/dask/common/base.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) 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, From eec940b15987508d981c94cc8602f16319f22f04 Mon Sep 17 00:00:00 2001 From: viclafargue Date: Wed, 25 Sep 2024 10:55:30 +0200 Subject: [PATCH 2/2] adding test --- .../cuml/cuml/tests/dask/test_dask_serialization.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) 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_")