Demonstration of new multiprocessing.shared_memory module

demo_shared_multiproc.py

import multiprocessing as mp
import time
from multiprocessing.shared_memory import SharedMemory

import numpy as np

SHAPE = (10_000_000,)
DTYPE = np.dtype("float32")


def _process_shm_at_idx(shm, idx, shape=SHAPE, dtype=DTYPE):
    a = np.ndarray(shape, dtype=dtype, buffer=shm.buf)
    a[idx] = 1
    # do other stuff with the array


def _process_npy_at_idx(arr, idx):
    arr[idx] = 1
    # do other stuff with the array


if __name__ == "__main__":

    t0 = time.time()
    in_arr = np.random.rand(*SHAPE).astype(DTYPE)

    shared_arr = SharedMemory(create=True, size=DTYPE.itemsize * SHAPE[0])
    b1 = np.ndarray(SHAPE, dtype=DTYPE, buffer=shared_arr.buf)
    b1[:] = in_arr
    print(f"Copied data into shared memory at {time.time() - t0:.2f} s")

    print(b1)
    _process_shm_at_idx(shared_arr, 0)
    # Should show 1 at index 0
    print(b1)

    procs = []
    with mp.Pool(20) as pool:
        for i in range(500):
            p = pool.apply_async(_process_shm_at_idx, (shared_arr, i))
            procs.append(p)
        for p in procs:
            p.get()

    print(b1)
    shared_arr.unlink()
    print(f"Done with shared memory version at {time.time() - t0:.2f} seconds")

    b2 = np.zeros(SHAPE, dtype=DTYPE)

    with mp.Pool(20) as pool:
        for i in range(500):
            p = pool.apply_async(_process_npy_at_idx, (b2, i))
            procs.append(p)

        for p in procs:
            p.get()

    print(f"Done with numpy version at {time.time() - t0:.2f} seconds")