Python multiprocessing throws Killed: 9

Question

I am trying to use multiprocessing to speed up a function where I tile 2000 arrays of shape (76, 76) into 3D arrays and apply a scaling factor.

It works fine when the number of tiles is less than about 200 but I get a Killed: 9 when it's greater than that and I need to be able to handle on order of 1000 tiles.

Here's a simplified version of the code:

from functools import partial
from multiprocessing.pool import ThreadPool
from multiprocessing import cpu_count
import numpy as np

def func_A(data, scale, N):
    """Tile the data N times and scale it"""
    arr = np.tile(data, (N, 1, 1))
    arr *= scale
    return arr

def func_B(N=4):
    """Create scaled arrays"""
    # Make data
    data = np.random.normal(size=(2000, 76, 76))

    # Make scales
    scales = np.arange(2000)

    # Multiprocess into tiled arrays
    pool = ThreadPool(cpu_count())
    func = partial(func_A, N=N)
    inpt = list(zip(data, scales))
    results = np.asarray(pool.starmap(func, inpt), dtype=np.float64)
    pool.close()
    pool.join()

    return results.swapaxes(0, 1)

So it's fine for func_B(4) but dies for func_B(500).

I understand I am taxing Python's memory with such large arrays but what is the best way to get func_B to work with large N... preferably quickly? Am I using multiprocessing wrong? Should I be using something else altogether, e.g. Dask, Numba, Cython, etc?

Any help would be greatly appreciated. Thanks!

Answer 1

I am not entirely sure what the purpose of your calculation is, but the following appears to do the job in dask

import dask.array as da
import numpy as np

# Make data
data = da.random.normal(size=(2000, 76, 76), chunks=(2000, 76, 76))

# Make scales
scales = np.arange(2000)
N = 500
out = da.repeat(data, N, axis=0).reshape((N, 2000, 76, 76)) * scales.reshape((1, 2000, 1, 1))
out = out.sum(axis=0).compute()

Keeps working memory <~5GB and uses most of your cores.

Answer 2

I think that the most intuitive solution to override the memory problem is to work with float16 arrays. I try to rewrite all the process in more simple way (func_C)

### your method ###

def func_A(data, scale, N):
    """Tile the data N times and scale it"""
    arr = np.tile(data, (N, 1, 1))
    arr *= scale
    return arr

def func_B(N=4):
    """Create scaled arrays"""
    # Make data
    data = np.random.normal(size=(2000, 76, 76)).astype(np.float16) ###### set float16

    # Make scales
    scales = np.arange(2000).astype(np.float16) ###### set float16

    # Multiprocess into tiled arrays
    pool = ThreadPool(cpu_count())
    func = partial(func_A, N=N)
    inpt = list(zip(data, scales))
    results = np.asarray(pool.starmap(func, inpt), dtype=np.float16) ###### set float16
    pool.close()
    pool.join()

    return results.swapaxes(0, 1)

### alternative method ###

def func_C(N=4):

    scales = np.arange(2000).astype(np.float16)
    data = np.random.normal(size=(2000, 76, 76)).astype(np.float16)
    results = np.stack(N*[data*scales[:,None,None]])

    return results

CHECK RESULTS

np.random.seed(33)
a = func_B(10)
np.random.seed(33)
b = func_C(10)
(a == b).all() # ===> TRUE

CHECK PERFORMANCES