Repository: cupy/cupy
Pull request: Allow `MatDescriptor` to be pickle-able
Fixes https://github.com/cupy/cupy/issues/3061

Problem statement:
cupy.sparse.MatDescriptor is not pickleable
I am encountering a segfault when freeing a CuPy sparse `csr_matrix` which was computed from a Dask array. I have an intuition that Dask might not be the cause of this problem. From a simple inspection of properties like `__cuda_array_interface__` and the allocator on the underlying index/data arrays, I don't see any obvious indications of a double free.

I am also able to successfully pickle and unpickle a CuPy `csr_matrix` and was able to free the resulting unpickled matrix without a segfault. 

Also, the strangest part about this error is that it doesn't happen when the contents of `build_arr` is run sequentially in the calling function. That is, I only get this segfault when returning a Dask.Array backed by sparse CuPy arrays from a function and calling "compute" in the calling function. 

```
from dask.distributed import Client
from dask_cuda import LocalCUDACluster
import cupy as cp
import dask.array

cluster = LocalCUDACluster()
client = Client(cluster)

x = cp.sparse.random(1000, 10).astype(cp.float32).tocsr()
cp.cuda.Stream.null.synchronize()

def build_arr(x):
  f = client.scatter(x)
  ret = dask.array.from_delayed(f, shape=x.shape,
                                                   meta=cp.sparse.csr_matrix(cp.zeros(1), dtype=x.dtype))
  return ret

ret = build_arr(x)
x_ret = ret.compute()
print(str(x_ret))
del x_ret
del ret
```

Here's the output and exception:

```
>>> print(str(x_ret))
  (3, 0)	0.93203676
  (5, 3)	0.816514
  (22, 0)	0.7288292
  (23, 5)	0.1614296
  (56, 2)	0.10485178
  (76, 7)	0.74339426
  (100, 5)	0.58028
  (111, 2)	0.85776955
  (129, 5)	0.6123046
  (133, 4)	0.60502696
  (140, 4)	0.74886125
  (151, 9)	0.315743
  (165, 2)	0.2791235
  (182, 5)	0.9034685
  (183, 4)	0.5153093
  (202, 3)	0.08361749
  (206, 1)	0.8262547
  (221, 0)	0.743914
  (223, 0)	0.6860539
  (252, 4)	0.15654124
  (256, 8)	0.4413018
  (274, 0)	0.72435445
  (274, 2)	0.56952196
  (286, 3)	0.17318754
  (290, 4)	0.1044073
  :	:
  (775, 2)	0.19066067
  (775, 5)	0.2838834
  (794, 2)	0.4538882
  (798, 6)	0.3432446
  (799, 2)	0.4626762
  (818, 1)	0.5111078
  (828, 4)	0.5325761
  (829, 8)	0.40714285
  (843, 2)	0.28924194
  (850, 3)	0.08357895
  (852, 9)	0.02581876
  (862, 8)	0.6738098
  (880, 0)	0.9008745
  (889, 5)	0.82182395
  (921, 1)	0.61768174
  (930, 1)	0.7116849
  (937, 3)	0.2526587
  (939, 6)	0.24009188
  (953, 6)	0.9483312
  (961, 2)	0.38405782
  (972, 0)	0.9323737
  (978, 9)	0.74389035
  (980, 0)	0.9883633
  (992, 1)	0.5483194
  (996, 3)	0.033570997
[deeplearn:11490:0:11493] Caught signal 11 (Segmentation fault: address not mapped to object at address 0x560443a7e7b8)
==== backtrace ====
    0  /share/conda/cuml_12_1/lib/python3.7/site-packages/ucp/_libs/../../../../libucs.so.0(+0x22631) [0x7fa90c643631]
    1  /share/conda/cuml_12_1/lib/python3.7/site-packages/ucp/_libs/../../../../libucs.so.0(+0x22802) [0x7fa90c643802]
    2  /lib64/libpthread.so.0(+0xf5f0) [0x7fa92064f5f0]
    3  /lib64/libc.so.6(cfree+0x1c) [0x7fa9202f7ecc]
    4  /share/conda/cuml_12_1/lib/python3.7/site-packages/cupy/cuda/../../../../libcusparse.so.10.0(cusparseDestroyMatDescr+0xe) [0x7fa89b971fde]
    5  /share/conda/cuml_12_1/lib/python3.7/site-packages/cupy/cuda/cusparse.cpython-37m-x86_64-linux-gnu.so(+0x6f7c1) [0x7fa8dc7867c1]
    6  /share/conda/cuml_12_1/bin/python(_PyMethodDef_RawFastCallKeywords+0x8d) [0x55dd1574382d]
    7  /share/conda/cuml_12_1/bin/python(_PyCFunction_FastCallKeywords+0x21) [0x55dd15743b21]
    8  /share/conda/cuml_12_1/bin/python(_PyEval_EvalFrameDefault+0x501e) [0x55dd157afd5e]
    9  /share/conda/cuml_12_1/bin/python(_PyEval_EvalCodeWithName+0x2f9) [0x55dd156ef729]
   10  /share/conda/cuml_12_1/bin/python(_PyFunction_FastCallDict+0x1d5) [0x55dd156f0865]
   11  /share/conda/cuml_12_1/bin/python(+0x182fe5) [0x55dd1575dfe5]
   12  /share/conda/cuml_12_1/bin/python(PyObject_CallFinalizer+0x85) [0x55dd15780735]
   13  /share/conda/cuml_12_1/bin/python(PyObject_CallFinalizerFromDealloc+0x1b) [0x55dd1578082b]
   14  /share/conda/cuml_12_1/bin/python(+0x1a5c7a) [0x55dd15780c7a]
   15  /share/conda/cuml_12_1/bin/python(+0x10d78c) [0x55dd156e878c]
   16  /share/conda/cuml_12_1/bin/python(+0x1a59e0) [0x55dd157809e0]
   17  /share/conda/cuml_12_1/bin/python(_PyDict_DelItem_KnownHash+0x36b) [0x55dd1575708b]
   18  /share/conda/cuml_12_1/bin/python(_PyEval_EvalFrameDefault+0x255a) [0x55dd157ad29a]
   19  /share/conda/cuml_12_1/bin/python(_PyFunction_FastCallDict+0x10b) [0x55dd156f079b]
   20  /share/conda/cuml_12_1/bin/python(_PyObject_FastCall_Prepend+0x65) [0x55dd15710275]
   21  /share/conda/cuml_12_1/bin/python(+0x17ffab) [0x55dd1575afab]
   22  /share/conda/cuml_12_1/bin/python(_PyEval_EvalFrameDefault+0x255a) [0x55dd157ad29a]
   23  /share/conda/cuml_12_1/bin/python(_PyFunction_FastCallDict+0x10b) [0x55dd156f079b]
   24  /share/conda/cuml_12_1/bin/python(_PyObject_FastCall_Prepend+0x65) [0x55dd15710275]
   25  /share/conda/cuml_12_1/bin/python(+0x17ffab) [0x55dd1575afab]
   26  /share/conda/cuml_12_1/bin/python(_PyEval_EvalFrameDefault+0x255a) [0x55dd157ad29a]
   27  /share/conda/cuml_12_1/bin/python(_PyFunction_FastCallDict+0x10b) [0x55dd156f079b]
   28  /share/conda/cuml_12_1/bin/python(_PyObject_FastCall_Prepend+0x65) [0x55dd15710275]
   29  /share/conda/cuml_12_1/bin/python(+0x17ffab) [0x55dd1575afab]
   30  /share/conda/cuml_12_1/bin/python(_PyEval_EvalFrameDefault+0x255a) [0x55dd157ad29a]
   31  /share/conda/cuml_12_1/bin/python(_PyFunction_FastCallDict+0x10b) [0x55dd156f079b]
   32  /share/conda/cuml_12_1/bin/python(_PyObject_FastCall_Prepend+0x65) [0x55dd15710275]
   33  /share/conda/cuml_12_1/bin/python(+0x17ffab) [0x55dd1575afab]
   34  /share/conda/cuml_12_1/bin/python(_PyEval_EvalFrameDefault+0x255a) [0x55dd157ad29a]
   35  /share/conda/cuml_12_1/bin/python(_PyFunction_FastCallDict+0x10b) [0x55dd156f079b]
   36  /share/conda/cuml_12_1/bin/python(_PyObject_FastCall_Prepend+0x65) [0x55dd15710275]
   37  /share/conda/cuml_12_1/bin/python(+0x17ffab) [0x55dd1575afab]
   38  /share/conda/cuml_12_1/bin/python(_PyEval_EvalFrameDefault+0x255a) [0x55dd157ad29a]
   39  /share/conda/cuml_12_1/bin/python(_PyEval_EvalCodeWithName+0x2f9) [0x55dd156ef729]
   40  /share/conda/cuml_12_1/bin/python(_PyFunction_FastCallKeywords+0x325) [0x55dd157431a5]
   41  /share/conda/cuml_12_1/bin/python(_PyEval_EvalFrameDefault+0x6a0) [0x55dd157ab3e0]
   42  /share/conda/cuml_12_1/bin/python(_PyGen_Send+0x2a2) [0x55dd1575c4f2]
   43  /share/conda/cuml_12_1/lib/python3.7/lib-dynload/_asyncio.cpython-37m-x86_64-linux-gnu.so(+0xc33e) [0x7fa91864633e]
   44  /share/conda/cuml_12_1/bin/python(_PyObject_FastCallKeywords+0x49b) [0x55dd1575b85b]
   45  /share/conda/cuml_12_1/bin/python(+0x2097c3) [0x55dd157e47c3]
   46  /share/conda/cuml_12_1/bin/python(_PyMethodDef_RawFastCallDict+0x194) [0x55dd15711a44]
   47  /share/conda/cuml_12_1/bin/python(_PyCFunction_FastCallDict+0x21) [0x55dd15711c81]
   48  /share/conda/cuml_12_1/bin/python(_PyEval_EvalFrameDefault+0x5da7) [0x55dd157b0ae7]
   49  /share/conda/cuml_12_1/bin/python(_PyFunction_FastCallKeywords+0xfb) [0x55dd15742f7b]
   50  /share/conda/cuml_12_1/bin/python(_PyEval_EvalFrameDefault+0x6a0) [0x55dd157ab3e0]
   51  /share/conda/cuml_12_1/bin/python(_PyFunction_FastCallKeywords+0xfb) [0x55dd15742f7b]
   52  /share/conda/cuml_12_1/bin/python(_PyEval_EvalFrameDefault+0x6a0) [0x55dd157ab3e0]
   53  /share/conda/cuml_12_1/bin/python(_PyFunction_FastCallKeywords+0xfb) [0x55dd15742f7b]
   54  /share/conda/cuml_12_1/bin/python(_PyEval_EvalFrameDefault+0x6a0) [0x55dd157ab3e0]
   55  /share/conda/cuml_12_1/bin/python(_PyFunction_FastCallKeywords+0xfb) [0x55dd15742f7b]
   56  /share/conda/cuml_12_1/bin/python(_PyEval_EvalFrameDefault+0x6a0) [0x55dd157ab3e0]
   57  /share/conda/cuml_12_1/bin/python(_PyEval_EvalCodeWithName+0xac9) [0x55dd156efef9]
   58  /share/conda/cuml_12_1/bin/python(_PyFunction_FastCallKeywords+0x387) [0x55dd15743207]
   59  /share/conda/cuml_12_1/bin/python(_PyEval_EvalFrameDefault+0x6a0) [0x55dd157ab3e0]
   60  /share/conda/cuml_12_1/bin/python(_PyEval_EvalCodeWithName+0xac9) [0x55dd156efef9]
   61  /share/conda/cuml_12_1/bin/python(_PyFunction_FastCallDict+0x400) [0x55dd156f0a90]
===================

```

What I find really strange about this error is this: `libcusparse.so.10.0(cusparseDestroyMatDescr+0xe)` 

I'm not sure completely sure where a cuSparse matrix descriptor would be lingering in the code. I also tried adding more stream synchronization just to make absolutely sure there wasn't a stray conversion somewhere, but that didn't seem to help.

Hints from issue discussion:
I believe I just figured out why this is happening:

```
>>> import cupy as cp
>>> import pickle
>>> 
>>> a = cp.sparse.random(1000, 100, format='csr', dtype=cp.float32)
>>> b = pickle.dumps(a)
>>> c = pickle.loads(b)
>>> c._descr.descriptor
94834805355424
>>> a._descr.descriptor
94834805355424

```
We should make this pickleable (eg. create a new mat descriptor each time, just like when a new sparse array is created). 
Update: I was able to fix this on my end using `copyreg`, though I think it would be worthwhile to fix this in CuPy as well. 

```
import cupy as cp
import copyreg


def serialize_mat_descriptor(m):
    return cp.cupy.cusparse.MatDescriptor.create, ()

copyreg.pickle(cp.cupy.cusparse.MatDescriptor, serialize_mat_descriptor)

```
Added PR ( https://github.com/cupy/cupy/pull/3157 ) to provide `MatDescriptor` a `__reduce__` method, which should fix this.

Commit message: Add a few tests for `MatDescriptor` directly

Patch to review:
diff --git a/cupy/cusparse.py b/cupy/cusparse.py
index a58ef0e4e89..e46f7e37d90 100644
--- a/cupy/cusparse.py
+++ b/cupy/cusparse.py
@@ -20,6 +20,9 @@ def create(cls):
         descr = cusparse.createMatDescr()
         return MatDescriptor(descr)
 
+    def __reduce__(self):
+        return self.create, ()
+
     def __del__(self, is_shutting_down=util.is_shutting_down):
         if is_shutting_down():
             return
diff --git a/tests/cupy_tests/test_cusparse.py b/tests/cupy_tests/test_cusparse.py
index 784ed1101ba..10e2754cc17 100644
--- a/tests/cupy_tests/test_cusparse.py
+++ b/tests/cupy_tests/test_cusparse.py
@@ -1,3 +1,4 @@
+import pickle
 import unittest
 
 import numpy
@@ -8,9 +9,23 @@
 
 import cupy
 from cupy import testing
+from cupy import cusparse
 from cupyx.scipy import sparse
 
 
+class TestMatDescriptor(unittest.TestCase):
+
+    def test_create(self):
+        md = cusparse.MatDescriptor.create()
+        assert isinstance(md.descriptor, int)
+
+    def test_pickle(self):
+        md = cusparse.MatDescriptor.create()
+        md2 = pickle.loads(pickle.dumps(md))
+        assert isinstance(md2.descriptor, int)
+        assert md.descriptor != md2.descriptor
+
+
 @testing.parameterize(*testing.product({
     'dtype': [numpy.float32, numpy.float64],
     'transa': [True, False],
diff --git a/tests/cupyx_tests/scipy_tests/sparse_tests/test_coo.py b/tests/cupyx_tests/scipy_tests/sparse_tests/test_coo.py
index ce52eefac87..1501201edc9 100644
--- a/tests/cupyx_tests/scipy_tests/sparse_tests/test_coo.py
+++ b/tests/cupyx_tests/scipy_tests/sparse_tests/test_coo.py
@@ -1,3 +1,4 @@
+import pickle
 import unittest
 
 import numpy
@@ -155,6 +156,14 @@ def test_init_copy_other_scipy_sparse(self, xp, sp):
         assert n.shape == m.shape
         return n
 
+    def test_pickle_roundtrip(self, xp, sp):
+        s = _make(xp, sp, self.dtype)
+        s2 = pickle.loads(pickle.dumps(s))
+        assert s._descr.descriptor != s2._descr.descriptor
+        assert s.shape == s2.shape
+        assert s.dtype == s2.dtype
+        assert (s != s2).count_nonzero()
+
     def test_shape(self):
         self.assertEqual(self.m.shape, (3, 4))
 
diff --git a/tests/cupyx_tests/scipy_tests/sparse_tests/test_csc.py b/tests/cupyx_tests/scipy_tests/sparse_tests/test_csc.py
index bce7887b222..5988f9cdb03 100644
--- a/tests/cupyx_tests/scipy_tests/sparse_tests/test_csc.py
+++ b/tests/cupyx_tests/scipy_tests/sparse_tests/test_csc.py
@@ -1,3 +1,4 @@
+import pickle
 import unittest
 
 import numpy
@@ -257,6 +258,14 @@ def test_toarray(self):
         self.assertTrue(m.flags.c_contiguous)
         cupy.testing.assert_allclose(m, expect)
 
+    def test_pickle_roundtrip(self, xp, sp):
+        s = _make(xp, sp, self.dtype)
+        s2 = pickle.loads(pickle.dumps(s))
+        assert s._descr.descriptor != s2._descr.descriptor
+        assert s.shape == s2.shape
+        assert s.dtype == s2.dtype
+        assert (s != s2).count_nonzero()
+
 
 @testing.parameterize(*testing.product({
     'dtype': [numpy.float32, numpy.float64, numpy.complex64, numpy.complex128],
diff --git a/tests/cupyx_tests/scipy_tests/sparse_tests/test_csr.py b/tests/cupyx_tests/scipy_tests/sparse_tests/test_csr.py
index b2903dea58c..e5e2d6d6920 100644
--- a/tests/cupyx_tests/scipy_tests/sparse_tests/test_csr.py
+++ b/tests/cupyx_tests/scipy_tests/sparse_tests/test_csr.py
@@ -1,3 +1,4 @@
+import pickle
 import unittest
 
 import numpy
@@ -270,6 +271,14 @@ def test_toarray(self):
         self.assertTrue(m.flags.c_contiguous)
         cupy.testing.assert_allclose(m, expect)
 
+    def test_pickle_roundtrip(self, xp, sp):
+        s = _make(xp, sp, self.dtype)
+        s2 = pickle.loads(pickle.dumps(s))
+        assert s._descr.descriptor != s2._descr.descriptor
+        assert s.shape == s2.shape
+        assert s.dtype == s2.dtype
+        assert (s != s2).count_nonzero()
+
 
 @testing.parameterize(*testing.product({
     'dtype': [numpy.float32, numpy.float64, numpy.complex64, numpy.complex128],
diff --git a/tests/cupyx_tests/scipy_tests/sparse_tests/test_dia.py b/tests/cupyx_tests/scipy_tests/sparse_tests/test_dia.py
index 1f1f773a188..7a58e05cc13 100644
--- a/tests/cupyx_tests/scipy_tests/sparse_tests/test_dia.py
+++ b/tests/cupyx_tests/scipy_tests/sparse_tests/test_dia.py
@@ -1,3 +1,4 @@
+import pickle
 import unittest
 
 
@@ -102,6 +103,14 @@ def test_toarray(self):
         self.assertTrue(m.flags.c_contiguous)
         cupy.testing.assert_allclose(m, expect)
 
+    def test_pickle_roundtrip(self, xp, sp):
+        s = _make(xp, sp, self.dtype)
+        s2 = pickle.loads(pickle.dumps(s))
+        assert s._descr.descriptor != s2._descr.descriptor
+        assert s.shape == s2.shape
+        assert s.dtype == s2.dtype
+        assert (s != s2).count_nonzero()
+
     def test_diagonal(self):
         testing.assert_array_equal(
             self.m.diagonal(-2), cupy.array([0], self.dtype))


Code review comments:
diff_hunk: @@ -155,6 +156,14 @@ def test_init_copy_other_scipy_sparse(self, xp, sp):
         assert n.shape == m.shape
         return n
 
+    def test_pickle_roundtrip(self, xp, sp):
+        s = _make(xp, sp, self.dtype)
path: tests/cupyx_tests/scipy_tests/sparse_tests/test_coo.py
text: @user1:
I think you just need this? 
```suggestion
    def test_pickle_roundtrip(self):
        s = _make(cupy, sparse, self.dtype)
```
(Likewise for all roundtrip tests.)

@author:
Huh interesting. I thought this would be tested with NumPy/SciPy first and then CuPy second. Is there a way for me to get that behavior here?

@user1:
You can do a quick loop like this: https://github.com/cupy/cupy/blob/7c96b16fe26277f383ebde190cc36d5518940565/tests/cupyx_tests/scipy_tests/sparse_tests/test_coo.py#L491-L492
But what's the point? Don't you just wanna test if CuPy sparse matrices are picklable? You can't compare CuPy's pickle with SciPy's, after all. 

@author:
Good point.

@author:
Fixed in the latest commit 🙂