[Python] Can `spawn` be made as memory efficient as `fork` with multiprocessing?

I am on Linux and have working multiprocessing code that uses fork. Here is a MWE version:

from multiprocessing import Pool
from time import perf_counter as now
import numpy as np


def make_func():
n = 20000
np.random.seed(6)
start = now()
M = np.random.rand(n, n)
return lambda x, y: M[x, x] + M[y, y]


class ParallelProcessor:
def __init__(self):
pass

def process_task(self, args):
"""Unpack arguments internally"""
index, integer_arg = args
print(f(index, integer_arg))

def run_parallel(self, tasks, num_cores=None):
"""Simplified parallel execution without partial"""
num_cores = num_cores

task_args = [(idx, val) for idx, val in enumerate(tasks)]
start = now()
global f
f = make_func()
print(f"************** {now() - start} seconds to make f")
start = now()
with Pool(num_cores) as pool:
results = pool.map( self.process_task, task_args)
print(f"************** {now() - start} seconds to run all jobs")
return results


if __name__ == "__main__":
processor = ParallelProcessor()
processor.run_parallel(tasks=[1, 2, 3, 4, 5, 6], num_cores=6)


This uses the fact that global f is shared with the workers and as the numpy array is not modified by them, it is not copied.

In 3.14 Linux multiprocessing will be moved to spawn. So in anticipation I made a version using spawn.

from multiprocessing import Pool, RawArray, set_start_method
from time import perf_counter as now
import numpy as np

def init_worker(shared_array_base, shape):
"""Initializer function to set up shared memory for each worker"""
global M
M = np.frombuffer(shared_array_base, dtype=np.float64).reshape(shape)

def worker_task(args):
"""Worker function that reconstructs f using shared memory"""
index, integer_arg = args
result = M[index, index] + M[integer_arg, integer_arg]
print(result)
return result

class ParallelProcessor:
def __init__(self):
pass

def run_parallel(self, tasks, num_cores=None):
"""Run tasks in parallel using spawn and shared memory"""
set_start_method('spawn', force=True) # Ensure 'spawn' is used

n = 20000
shape = (n, n)

# Use 'd' for double-precision float (float64) instead of np.float64
shared_array_base = RawArray('d', n * n)
M_local = np.frombuffer(shared_array_base, dtype=np.float64).reshape(shape)

# Initialize the array in the main process
np.random.seed(7)
start = now()
M_local[:] = np.random.rand(n, n)
print(f"************** {now() - start} seconds to make M")

# Prepare arguments for worker tasks
task_args = [(idx, val) for idx, val in enumerate(tasks)]

start = now()
with Pool(num_cores, initializer=init_worker, initargs=(shared_array_base, shape)) as pool:
results = pool.map(worker_task, task_args)
print(f"************** {now() - start} seconds to run all jobs")
return results

if __name__ == "__main__":
processor = ParallelProcessor()
processor.run_parallel(tasks=[1, 2, 3, 4, 5, 6], num_cores=6)


Unfortunately, not only is this slower, it also seems to use much more memory. Maybe it is making a copy of the numpy array?

Is there a way to make the spawn version as efficient as the fork version?

Continue reading...