Benchmarking a data-intensive operation - ATM 622: General Circulation of the Atmosphere

We’re going to compute the zonal, time average of a product of two high-frequency fields from the CFSR dataset.

Specifically we will compute

[\overline{vT}]

(1)

for the month of January (one single year)

I want to compare various ways of accelerating the calculation.

Compute environment¶

I logged into https://jupyterlab.arcc.albany.edu, and spawned on batch with 4 cores and 32 GB memory

This resource is open to anyone with UAlbany login credentials.

Python kernel is daes_nov22.

Eager execution in memory¶

First we do the “normal” style: open the datasets, execute the calculation immediately

import xarray as xr

files = ['/network/daes/cfsr/data/2021/v.2021.0p5.anl.nc',
         '/network/daes/cfsr/data/2021/t.2021.0p5.anl.nc']

v_ds = xr.open_dataset(files[0])
T_ds = xr.open_dataset(files[1])
v_ds

Computing the product $vT$ over one month involves taking the product of two arrays of size

124 x 32 x 361 x 720

The total number of grid points for the whole year is

v_ds.time.size * v_ds.lev.size * v_ds.lat.size * v_ds.lon.size

12143462400

The month of January alone is about $1\times 10^9$ points.

Using 32-bit real numbers (4 bytes per numbers), this means we’re taking the product of two arrays of size around 4 GB each.

Timing the eager execution¶

%time (v_ds.v.groupby(v_ds.time.dt.month)[1] * T_ds.t.groupby(T_ds.time.dt.month)[1]).mean(dim=('lon','time'))

CPU times: user 38 s, sys: 6.03 s, total: 44 s
Wall time: 46.1 s

v_ds.close()
T_ds.close()

Lazy execution using Dask locally¶

ds = xr.open_mfdataset(files, chunks={'time':30, 'lev': 1})
ds

vT = ds.v * ds.t
meanfield = vT.mean(dim='lon').groupby(ds.time.dt.month).mean()
meanfield

%time meanfield.sel(month=1).load()

CPU times: user 49.3 s, sys: 3.27 s, total: 52.5 s
Wall time: 43.1 s

ds.close()

Conclusion¶

On this particular system, we seem to get about the same performance either way -- giving advantage to the lazy method, because the code is nicer to work with.

Lazy execution using Dask and a distributed cluster¶

Now we’ll try the same thing but farm the computation out to a dask cluster.

from dask_jobqueue import SLURMCluster
from dask.distributed import Client

cluster = SLURMCluster(processes=4,  #By default Dask will run one Python process per job. 
                       # However, you can optionally choose to cut up that job into multiple processes
                       cores=80, # size of a single job -- typically one node of the HPC cluster
                       memory="376.2GB",  # the max memory on the 80-core batch nodes, I believe
                       walltime="01:00:00",
                       queue="batch",
                      )
cluster.scale(1)
client = Client(cluster)
client

Effects of different chunking¶

ds = xr.open_mfdataset(files, chunks={'time':30, 'lev': 1})
vT = ds.v * ds.t
vT

meanfield = vT.mean(dim='lon').groupby(ds.time.dt.month).mean()
meanfield

%time meanfield.sel(month=1).load()

CPU times: user 994 ms, sys: 224 ms, total: 1.22 s
Wall time: 18.4 s

ds.close()

Same thing, different chunking¶

ds = xr.open_mfdataset(files, chunks={'time':10, 'lev': 32})

meanfield = (ds.v*ds.t).mean(dim='lon').groupby(ds.time.dt.month).mean()
meanfield

%time meanfield.sel(month=1).load()

CPU times: user 261 ms, sys: 62.7 ms, total: 324 ms
Wall time: 20.2 s

ds.close()

And yet another chunking¶

ds = xr.open_mfdataset(files, chunks={'time':124, 'lev': 32})
vT = ds.v * ds.t
vT

meanfield = vT.mean(dim='lon').groupby(ds.time.dt.month).mean()
meanfield

%time meanfield.sel(month=1).load()

CPU times: user 475 ms, sys: 119 ms, total: 593 ms
Wall time: 45 s

ds.close()

Results¶

The smallest wall time for this calculation was achieved using the distributed cluster and a chunking strategy chunks={'time':10, 'lev': 32}.

For this one-month calculation, using the distributed cluster reduces our wall time from about 50 s down to about 20 s.

That’s close to a factor of three speed-up.

Can we do better through more intelligent chunking and/or fine-tuning of the distributed cluster? Probably.

cluster.close()

2022-11-17 22:49:54,749 - distributed.client - ERROR - Failed to reconnect to scheduler after 30.00 seconds, closing client