update tests for cloudpickle

Author: Spaak

pymc3/tests/test_distributions.py

@@ -3109,9 +3109,9 @@ def func(x):
         y = pm.DensityDist("y", func)
         pm.sample(draws=5, tune=1, mp_ctx="spawn")
 
-    import pickle
+    import cloudpickle
 
-    pickle.loads(pickle.dumps(y))
+    cloudpickle.loads(cloudpickle.dumps(y))
 
 
 def test_distinct_rvs():

pymc3/tests/test_minibatches.py

@@ -13,9 +13,9 @@
 #   limitations under the License.
 
 import itertools
-import pickle
 
 import aesara
+import cloudpickle
 import numpy as np
 import pytest
 
@@ -132,10 +132,10 @@ def gen():
 
     def test_pickling(self, datagen):
         gen = generator(datagen)
-        pickle.loads(pickle.dumps(gen))
+        cloudpickle.loads(cloudpickle.dumps(gen))
         bad_gen = generator(integers())
-        with pytest.raises(Exception):
-            pickle.dumps(bad_gen)
+        with pytest.raises(TypeError, match="cannot pickle 'generator' object"):
+            cloudpickle.dumps(bad_gen)
 
     def test_gen_cloning_with_shape_change(self, datagen):
         gen = generator(datagen)

pymc3/tests/test_model.py

@@ -11,14 +11,14 @@
 #   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 #   See the License for the specific language governing permissions and
 #   limitations under the License.
-import pickle
 import unittest
 
 from functools import reduce
 
 import aesara
 import aesara.sparse as sparse
 import aesara.tensor as at
+import cloudpickle
 import numpy as np
 import numpy.ma as ma
 import numpy.testing as npt
@@ -407,9 +407,7 @@ def test_model_pickle(tmpdir):
         x = pm.Normal("x")
         pm.Normal("y", observed=1)
 
-    file_path = tmpdir.join("model.p")
-    with open(file_path, "wb") as buff:
-        pickle.dump(model, buff)
+    cloudpickle.loads(cloudpickle.dumps(model))
 
 
 def test_model_pickle_deterministic(tmpdir):
@@ -420,9 +418,7 @@ def test_model_pickle_deterministic(tmpdir):
         pm.Deterministic("w", x / z)
         pm.Normal("y", observed=1)
 
-    file_path = tmpdir.join("model.p")
-    with open(file_path, "wb") as buff:
-        pickle.dump(model, buff)
+    cloudpickle.loads(cloudpickle.dumps(model))
 
 
 def test_model_vars():

pymc3/tests/test_parallel_sampling.py

@@ -71,12 +71,6 @@ def _crash_remote_process(a, master_pid):
     return 2 * np.array(a)
 
 
-def test_dill():
-    with pm.Model():
-        pm.Normal("x")
-        pm.sample(tune=1, draws=1, chains=2, cores=2, pickle_backend="dill", mp_ctx="spawn")
-
-
 def test_remote_pipe_closed():
     master_pid = os.getpid()
     with pm.Model():
@@ -112,7 +106,6 @@ def test_abort():
             mp_ctx=ctx,
             start={"a": np.array([1.0]), "b_log__": np.array(2.0)},
             step_method_pickled=None,
-            pickle_backend="pickle",
         )
         proc.start()
         while True:
@@ -147,7 +140,6 @@ def test_explicit_sample():
         mp_ctx=ctx,
         start={"a": np.array([1.0]), "b_log__": np.array(2.0)},
         step_method_pickled=None,
-        pickle_backend="pickle",
     )
     proc.start()
     while True:

pymc3/tests/test_pickling.py

@@ -15,6 +15,8 @@
 import pickle
 import traceback
 
+import cloudpickle
+
 from pymc3.tests.models import simple_model
 
 
@@ -26,8 +28,8 @@ def test_model_roundtrip(self):
         m = self.model
         for proto in range(pickle.HIGHEST_PROTOCOL + 1):
             try:
-                s = pickle.dumps(m, proto)
-                pickle.loads(s)
+                s = cloudpickle.dumps(m, proto)
+                cloudpickle.loads(s)
             except Exception:
                 raise AssertionError(
                     "Exception while trying roundtrip with pickle protocol %d:\n" % proto

pymc3/tests/test_variational_inference.py

@@ -757,7 +757,7 @@ def test_remove_scan_op():
 
 
 def test_clear_cache():
-    import pickle
+    import cloudpickle
 
     with pm.Model():
         pm.Normal("n", 0, 1)
@@ -767,7 +767,7 @@ def test_clear_cache():
         inference.approx._cache.clear()
         # should not be cleared at this call
         assert all(len(c) == 0 for c in inference.approx._cache.values())
-        new_a = pickle.loads(pickle.dumps(inference.approx))
+        new_a = cloudpickle.loads(cloudpickle.dumps(inference.approx))
         assert not hasattr(new_a, "_cache")
         inference_new = pm.KLqp(new_a)
         inference_new.fit(n=10)
@@ -871,26 +871,26 @@ def test_rowwise_approx(three_var_model, parametric_grouped_approxes):
 
 
 def test_pickle_approx(three_var_approx):
-    import pickle
+    import cloudpickle
 
-    dump = pickle.dumps(three_var_approx)
-    new = pickle.loads(dump)
+    dump = cloudpickle.dumps(three_var_approx)
+    new = cloudpickle.loads(dump)
     assert new.sample(1)
 
 
 def test_pickle_single_group(three_var_approx_single_group_mf):
-    import pickle
+    import cloudpickle
 
-    dump = pickle.dumps(three_var_approx_single_group_mf)
-    new = pickle.loads(dump)
+    dump = cloudpickle.dumps(three_var_approx_single_group_mf)
+    new = cloudpickle.loads(dump)
     assert new.sample(1)
 
 
 def test_pickle_approx_aevb(three_var_aevb_approx):
-    import pickle
+    import cloudpickle
 
-    dump = pickle.dumps(three_var_aevb_approx)
-    new = pickle.loads(dump)
+    dump = cloudpickle.dumps(three_var_aevb_approx)
+    new = cloudpickle.loads(dump)
     assert new.sample(1000)