Benchmarking a data-intensive operation (snow version) - 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 snow with 8 cores and 64 GB memory

This resource is only open to snow group members

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 56.9 s, sys: 7.19 s, total: 1min 4s
Wall time: 1min 8s

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 1min 12s, sys: 5.86 s, total: 1min 18s
Wall time: 1min 5s

ds.close()

Conclusion¶

Similar to running on batch, we get about the same performance either way -- giving advantage to the lazy method, because the code is nicer to work with.

Snow is slower than batch -- probably just because the cpu hardware is significantly older.

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=8,  #By default Dask will run one Python process per job. 
                       # However, you can optionally choose to cut up that job into multiple processes
                       cores=32,  # size of a single job -- typically one node of the HPC cluster
                       memory="256GB", # snow has 8 GB ram per physical core I believe
                       walltime="01:00:00",
                       queue="snow",
                       interface="ib0",
                      )
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 2.5 s, sys: 455 ms, total: 2.96 s
Wall time: 41.9 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 812 ms, sys: 120 ms, total: 932 ms
Wall time: 29.9 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 1.7 s, sys: 217 ms, total: 1.91 s
Wall time: 1min 13s

ds.close()

Results¶

Same as for batch, 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 65 s down to about 30 s.

That’s about a factor of two speedup.

Compared to batch, snow is about 1.5x slower (30 seconds vs 20 seconds).

cluster.close()