I have a Python class that I use for a series of computations in which I am initialising a number of different pandas.DataFrames stored inside nested dictionaries. The bottleneck in initialising the class is as follows:
I have a regular grid of points spanning the Earth and a set of polygons (representing tectonic plates) that also cover the Earth. Of each point in the grid, I need to find in which plate they are located and assign the ID of that plate to a list.
Because this is relatively slow (i.e. takes 10s of seconds per iteration instead of <1s), I would like to parallelise this section (and only this section!) of the class. I am new to parallelisation, however, and I ran into some issues.
Here is the function that handles finding the plateID.
def process_plateIDs(
geometries_data: list,
lats_chunk: _numpy.array,
lons_chunk: _numpy.array,
) -> list:
plateIDs = _numpy.zeros(len(lats_chunk))
for topology_geometry, topology_plateID in geometries_data:
mask = topology_geometry.contains(_geopandas.points_from_xy(lons_chunk, lats_chunk))
plateIDs[mask] = topology_plateID
return plateIDs
Where I simplified the plate geometries that are originally stored in a geopandas.DataFrame to reduce the computation time:
def extract_geometry_data(topology_geometries):
return [(geom, plateID) for geom, plateID in zip(topology_geometries.geometry, topology_geometries.PLATEID1)]
Note that each instance topology_geometries is stored in attribute of the main class self.topology_geometries, which is a dictionary that contains the geometries of the tectonic plates for different ages and cases (so self.topology_geometries[age][case] gives one geopandas.DataFrame).
I am trying to run this section in parallel by feeding it to Pool from multiprocessing:
# Use all available CPUs if num_workers is not specified
if num_workers is None:
num_workers = _os.cpu_count()
# Split the data into chunks
chunk_size = len(lats) // num_workers
chunks = [(geometries_data.copy(), lats[i:i + chunk_size].copy(), lons[i:i + chunk_size].copy()) for i in range(0, len(lats), chunk_size)]
# Create a Pool of workers
with Pool(num_workers) as pool:
# Map the process_chunk function to chunks
results = pool.starmap(process_plateIDs, chunks)
# Concatenate results from all chunks
plateIDs = _numpy.concatenate(results)
However, this results in each worker reinitialising the entire class, which is not at all what I intended. I suspect this is because I am passing a view instead of separate copies of certain variables, which I unsuccessfully tried to prevent that by making explicit copies of each variable (as in the above code snippet).
How can I tell Python to only explicitly run process_plateIDs parallel, and keep the rest of the code sequential? It should run on a jupyter notebook.
Many thanks!
EDIT:
Additional information: I am running this on MacOS with M1 chip.
thomas.s is a new contributor to this site. Take care in asking for clarification, commenting, and answering.
Check out our Code of Conduct.
4
MacOS uses spawn as a multiprocessing.set_start_method so you need your script to be like this, see Compulsory usage of if name==”main” in windows while using multiprocessing
import all_needed_modules
if __name__ == "__main__":
do_actual_code()
this will prevent workers from running the code when they are importing your script.
the second problem is passing huge data around, all arguments are serialized and de-serialized each function call, so calling a function with big data blocks each time is bad, instead use a global and an initializer to set it, this initializer will only happen once per worker.
some_global_data = None
def set_global_data(data):
global some_global_data
some_global_data = data
...
...
# much later, inside the "__main__" block
with Pool(num_workers, initializer=set_global_data, initargs=(geometries_data,)) as pool:
# Map the process_chunk function to chunks
results = pool.starmap(process_plateIDs, chunks)
and in your function just use this global instead of passing the huge data block each time.
lastly the problem of jupyter …. well, jupyter has very bad support for multiprocessing on any OS except linux due to the absence of fork, so you need the function that you are calling to be in another file that you import separately, see Jupyter notebook never finishes processing using multiprocessing
from some_file import process_plateIDs, set_global_data
# call pool.startmap(process_plateIDs, chunks) later
another alternative is to run jupyter through a linux container so python can use fork and you can declare the multiprocessed function directly within your notebook without the extra complications of jupyter.