How to pickle objects with many views of same cython memoryview

Question

I have a number of objects that I am trying to pickle which all share the same (large) cython memoryview as an attribute. Since the memoryviews are passed by reference, they all share the same memory and the implementation is memory efficient.

Now I need to pickle these objects and reload them while keeping the shared data shared (if the shared data becomes not shared then the file size blows up and it is impossible to read into memory). Normally I think pickle recognizes shared data and just pickles/unpickles it once, but because memory views can't be pickled directly, they need to be converted to a numpy array in the reduce method for each object and pickle no longer recognizes that the data is shared.

Is there some way that I can maintain the shared data through the pickle/unpickle process?

A MWE follows:

import numpy as np    
import pickle

cdef class SharedMemory:  
    cdef public double[:, :] data

    def __init__(self, data):   
        self.data = data

    def duplicate(self):
        return SharedMemory(self.data)

    def __reduce__(self):   
        return self.__class__, (np.asarray(self.data),)


def main():   
    x = SharedMemory(np.random.randn(100, 100))

    duplicates = [x.duplicate() for _ in range(5)]

    cdef double* pointerx = &x.data[0, 0]
    cdef double* pointerd
    cdef double[:, :] ddata

    for d in duplicates:
        ddata = d.data 
        pointerd = &ddata[0, 0]
        if pointerd != pointerx:
            print('Memory is not shared')
        else:
            print('Memory is shared')

    print('pickling')
    with open('./temp.pickle', 'wb') as pfile:
        pickle.dump(x, pfile, protocol=pickle.HIGHEST_PROTOCOL)
        for d in duplicates:
            pickle.dump(d, pfile, protocol=pickle.HIGHEST_PROTOCOL)

    with open('./temp.pickle', 'rb') as pfile:
        nx = pickle.load(pfile)
        nd = []
        for d in duplicates:
            nd.append(pickle.load(pfile))

    ddata = nx.data
    cdef double* pointernx = &ddata[0, 0]

    for d in nd:
        ddata = d.data
        pointerd = &ddata[0, 0]
        if pointerd != pointernx:
            print('Memory is not shared')
        else:
            print('Memory is shared')

Put the above in a file test.pyx an cythonize with "cythonize -a -i test.pyx". Then "export PYTHONPATH="$PYTHONPATH":." and run

from test import main
main()

from python.

user_5769463 · Answer

There are actually two problems:

First: The shared objects are also shared after dump/load only if they where pickled in one go (see also this answer).

That means you need to do the following (or similar):

...
with open('./temp.pickle', 'wb') as pfile:
     pickle.dump((x,duplicates), pfile, protocol=pickle.HIGHEST_PROTOCOL)
...
with open('./temp.pickle', 'rb') as pfile:
     nx, nd = pickle.load(pfile)
...

When you dump single objects, pickle isn't able to track identical objects - doing so would be an issue: objects with the same id between two dump-calls could be completely different objects or the same objects with different content!

Second: You should not create new objects, but pass the shared numpy-object in __reduce__ (pickle doesn't look inside of a numpy-array to see, whether buffer is shared or not, but only at the id of the array):

def __reduce__(self):
    return self.__class__, (self.data.base,)

which will give you the desired result. data.base is a reference to the underlying original numpy-array (or whatever type, which must support pickling/unpickling, obviously).

Warning: As @DavidW has rightly pointed out, additional considerations must be taken into account, when working with sliced memory-views - because in this case base might not be "the same" as the actual memory view.

How to pickle objects with many views of same cython memoryview

1 Answers