Speeding up certain MetPy functions with Python multiprocessing

[DanielAdriaansen]

I asked a question over on #1533 which is still an open issue. I don't feel comfortable tackling the problem described there in terms of the MetPy code base and architecture, so I came up with a workaround that works for me outside of MetPy.

This workaround uses the precipitable_water() function, but it may be applicable to other functions within MetPy as well (like cape_cin, or anything possibly in the documentation that says "only functions on 1D profiles"). The assumed output from this example is a 2D array. The data used here is ERA5 reanalysis data on pressure levels, subset to a 181 (latitude) x 281 (longitude) domain centered over the conterminous United States.

The general concept of my workaround is to define a 2D ndarray to store the output, break it up into tiles, and then write a function to wrap precipitable_water() so that it can be called on multiple tiles of the output 2D grid simultaneously using the multiprocessing Python package. Then the tiles are unpacked and put into the final 2D output array.

Here's the code (it assumes you have the required inputs to precipitable_water() stored in an Xarray Dataset named ds):

# Set the x and y dimension sizes
nx = ds.dims['latitude']
ny = ds.dims['longitude']

# Define the 2D ndarray to hold the output
pw = np.empty([nx,ny])

# First, set how many tiles we want in x and y direction
ntilesx = 4
ntilesy = 4

# Figure out the length of the x side of each tile in number of cells
xlength = round(nx/ntilesx)

# Figure out the length of the y side of each tile in number of cells
ylength = round(ny/ntilesy)

# Figure out the starting and ending x values of each tile
tile_start_x = np.arange(0,xlength*ntilesx-1,xlength)
tile_end_x = tile_start_x+(xlength-1)
tile_end_x[-1] = nx-1

# Figure out the starting and ending y values of each tile
tile_start_y = np.arange(0,ylength*ntilesy-1,ylength)
tile_end_y = tile_start_y+(ylength-1)
tile_end_y[-1] = ny-1

# We need to repeat the tile indicies, such that each start_x and end_x gets run for
# each pair of start_y and end_y. Thus, we need ntilesx X ntilesy number of pairs.
# np.repeat takes (0,1) and makes (0,0,1,1) where repeats is N
# np.tile takes (0,1) and makes (0,1,0,1) where tiles is N
xbeg = np.repeat(tile_start_x,ntilesx)
xend = np.repeat(tile_end_x,ntilesx)
ybeg = np.tile(tile_start_y,ntilesy)
yend = np.tile(tile_end_y,ntilesy)

# Define a function to compute pwat to use with multiprocessing. This will be called for each tile in parallel.
def pwat_on_tile(pw,xb,xe,yb,ye):
  for x in range(xb,xe+1,1):
    for y in range(yb,ye+1,1):
      pw[x,y] = mpcalc.precipitable_water(ds['P'].isel(latitude=x,longitude=y),ds['TD'].isel(latitude=x,longitude=y),bottom=pbot,top=ptop).m
  return(pw)

# Define a pool of workers
mp = multiprocessing.Pool(max(multiprocessing.cpu_count()-2,1))

# Execute the function on each tile and store the results.
results = mp.starmap(pwat_on_tile,[(pw,x1,x2,y1,y2) for x1,x2,y1,y2 in tuple(zip(xbeg,xend,ybeg,yend))])

# Unpack the results. Each item in the results is just the full ndarray, but with good data only
# on the tile defined by xbeg,xend,ybeg,yend. So we need to extract that tile and overwrite the data
# in the full array, thus filling out the array from all tiles.
for r,x1,x2,y1,y2 in tuple(zip(results,xbeg,xend,ybeg,yend)):
  pw[x1:x2+1,y1:y2+1] = r[x1:x2+1,y1:y2+1]

print(pw.min())
print(pw.max())

Here are some time output using the Linux time command:
ntilesx = 1
ntilesy = 1

real	5m50.047s
user	5m50.280s
sys	0m1.536s

ntilesx = 2
ntilesy = 2

real	1m32.418s
user	5m57.764s
sys	0m1.552s

ntilesx = 3
ntilesy = 3

real	0m43.374s
user	6m1.896s
sys	0m1.576s

ntilesx = 4
ntilesy = 4

real	0m33.868s
user	7m24.220s
sys	0m2.772s

After using 16 tiles (ntilesx=4, ntilesy=4, 16 total tiles) on this particular hardware, there's no more speedup and 30s is about as fast as I could get it to run. But this is miles ahead of the nearly 6 minutes when using the full domain (90%+ speedup!).

Some things to keep in mind:

1. The optimal number of tiles in the x/y dimension may be different for you depending on your hardware and grid size
2. I am not a multiprocessing expert and there may be better ways to do what I am demonstrating

Here's the output from the serial version using nested for loops on the full grid (or a single tile):

min=1.0289893631330012
max=45.26808312559193

And here's the output using the multiprocessing method demonstrated above using 16 tiles:

min=1.0289893631330012
max=45.26808312559193

I welcome any and all feedback or suggestions, and hope that this is useful to others.

[ShareRingNS]

i got a BUG

Traceback (most recent call last):
  File "<string>", line 1, in <module>
  File "C:\Users\anaconda3\envs\py38\lib\multiprocessing\spawn.py", line 116, in spawn_main
    exitcode = _main(fd, parent_sentinel)
  File "C:\Users\anaconda3\envs\py38\lib\multiprocessing\spawn.py", line 125, in _main
    prepare(preparation_data)
  File "C:\Users\anaconda3\envs\py38\lib\multiprocessing\spawn.py", line 236, in prepare
    _fixup_main_from_path(data['init_main_from_path'])
  File "C:\Users\anaconda3\envs\py38\lib\multiprocessing\spawn.py", line 287, in _fixup_main_from_path
    main_content = runpy.run_path(main_path,
  File "C:\Users\anaconda3\envs\py38\lib\runpy.py", line 265, in run_path
    return _run_module_code(code, init_globals, run_name,
  File "C:\Users\anaconda3\envs\py38\lib\runpy.py", line 97, in _run_module_code
    _run_code(code, mod_globals, init_globals,
  File "C:\Users\anaconda3\envs\py38\lib\runpy.py", line 87, in _run_code
    exec(code, run_globals)
  File "C:\Users\Desktop\nudt_work\exp_doc\Toolbox\Index_Calculation3.py", line 172, in <module>
    mp = multiprocessing.Pool(max(multiprocessing.cpu_count()-2,1))
  File "C:\Users\anaconda3\envs\py38\lib\multiprocessing\context.py", line 119, in Pool
    return Pool(processes, initializer, initargs, maxtasksperchild,
  File "C:\Users\anaconda3\envs\py38\lib\multiprocessing\pool.py", line 212, in __init__
    self._repopulate_pool()
  File "C:\Users\anaconda3\envs\py38\lib\multiprocessing\pool.py", line 303, in _repopulate_pool
    return self._repopulate_pool_static(self._ctx, self.Process,
  File "C:\Users\anaconda3\envs\py38\lib\multiprocessing\pool.py", line 326, in _repopulate_pool_static
    w.start()
  File "C:\Users\anaconda3\envs\py38\lib\multiprocessing\process.py", line 121, in start
    self._popen = self._Popen(self)
  File "C:\Users\anaconda3\envs\py38\lib\multiprocessing\context.py", line 327, in _Popen
    return Popen(process_obj)
  File "C:\Users\anaconda3\envs\py38\lib\multiprocessing\popen_spawn_win32.py", line 45, in __init__
    prep_data = spawn.get_preparation_data(process_obj._name)
  File "C:\Users\anaconda3\envs\py38\lib\multiprocessing\spawn.py", line 154, in get_preparation_data
    _check_not_importing_main()
  File "C:\Users\anaconda3\envs\py38\lib\multiprocessing\spawn.py", line 134, in _check_not_importing_main
    raise RuntimeError('''
RuntimeError: 
        An attempt has been made to start a new process before the
        current process has finished its bootstrapping phase.

        This probably means that you are not using fork to start your
        child processes and you have forgotten to use the proper idiom
        in the main module:

            if __name__ == '__main__':
                freeze_support()
                ...

        The "freeze_support()" line can be omitted if the program
        is not going to be frozen to produce an executable.

Process finished with exit code -1

How can i solve it?

[dopplershift]

I've never seen that error before, but it looks like it's completely related to how you're using multiprocessing. I don't see anything related to MetPy in that traceback.