How to properly parallelizing a for-loop in python with shared variables?

Question

This is my first attempt for multiprocessing in python and I am having hard time finding a fast solution for parallelizing my code. I do my best to explain my situation.

This is the code I am trying to parallelizing (the main problem is the for-loop):

Original code:

if __name__ == "__main__":
    n = 10
    filling_array = np.ones((n, n))*-1
    reading_array = np.ones((n, n))*10
    
    for i in range(n):
        for j in range(int(i/2)):
            filling_array[i, j] = 10*i + reading_array[i, j]    

    print(filling_array.sum(axis=1), 'end')

And this is my solution:

def child_p(i,
            filling_array,
            reading_array):

    for j in range(int(i/2)):
        filling_array[i, j] = 10*i + reading_array[i, j]
    
    return filling_array



if __name__ == "__main__":
    n = 10
    filling_array = np.ones((n, n))*-1
    reading_array = np.ones((n, n))*10

    final_array = np.zeros((n, n))

    with Pool(cpu_count()) as pool:
        result = pool.starmap(child_p,
                              zip(
                                 range(n),
                                 repeat(filling_array, n),
                                 repeat(reading_array, n)
                                 )
                              )
            
    for iter in range(n):
        final_array[iter, :] = result[iter][iter, :]

    print(final_array.sum(axis=1), 'end')

Literally, the child function act like the for-loop to fill the filling_array, using the reading array. I have separated the "for-loop" as a child function.

My solution is faster than the original for-loop but it is very slow and memory-heavy!

I think the main problem is where the pool.starmap() tries to copying/mapping the two arrays into the function (the arrays are very large).

Would you please guide me? Is there a way that I share these large arrays (in my original code, there are 5 reading_array and 1 filling_array and they are large)?

How can I do that?

Answer 1

This is what I did after, in order to solve my problem, it somehow works, but Im not sure how good it is:

Edited

I have tried the method in this website, but for two arrays, like this:

def child_p(i,
            filling_array_specs, 
            reading_array_specs,):

    fa_name, fa_shape, fa_dtype = filling_array_specs
    fa_shm = SharedMemory(fa_name)
    filling_array = np.ndarray(shape=fa_shape, dtype=fa_dtype, buffer=fa_shm.buf)
    
    ra_name, ra_shape, ra_dtype = reading_array_specs
    ra_shm = SharedMemory(ra_name)
    reading_array = np.ndarray(shape=ra_shape, dtype=ra_dtype, buffer=ra_shm.buf)
    
    for j in range(int(i/2)):
        filling_array[i, j] = 10*i + reading_array[i, j]


if __name__ == "__main__":
    n = 10
    filling_array = np.ones((n, n))*-1
    reading_array = np.ones((n, n))*10
    
    fa_shape, fa_dtype = filling_array.shape, filling_array.dtype
    ra_shape, ra_dtype = reading_array.shape, reading_array.dtype

    with SharedMemoryManager() as smm:
        fa_shm = smm.SharedMemory(filling_array.nbytes)
        shm_filling_array = np.ndarray(shape=fa_shape, dtype=fa_dtype, buffer=fa_shm.buf)
        np.copyto(shm_filling_array, filling_array)
        filling_array_specs = [fa_shm.name, fa_shape, fa_dtype]
        
        ra_shm = smm.SharedMemory(reading_array.nbytes)
        shm_reading_array = np.ndarray(shape=ra_shape, dtype=ra_dtype, buffer=ra_shm.buf)
        np.copyto(shm_reading_array, reading_array)
        reading_array_specs = [ra_shm.name, ra_shape, ra_dtype]
        
        with Pool(cpu_count()) as pool:
            pool.starmap(child_p,
                         zip(
                             range(n),
                             repeat(filling_array_specs, n),
                             repeat(reading_array_specs, n)
                             )
                         )
            
    print(shm_filling_array.sum(axis=1), 'end')

Please note that I passed each array's specs [sh_name, shape, dtype] to the child_p.

So far it seems working, but I am not sure how fast it is with my large arrays, will update this post regarding the time.

Edit 2:

# for n = 2000
# The original for-loop took 82.13 sec
# The shared memory method took 17.39 sec

Is this the fastest method I can achive? Or there are other methods that you think would works better?