ocean-transport / scale-aware-air-sea

Repo for collaborative project on scale-aware air-sea fluxes
2 stars 0 forks source link

Performance issues when using xESMF horizontal regridding in the cloud #12

Closed jbusecke closed 2 years ago

jbusecke commented 2 years ago

I have started to work on the CM2.6 data, and am running into issues with memory useage. When I try to create a regridder object for the full global grid

import xarray as xr
import intake
import xesmf as xe

from intake import open_catalog
kwargs = dict(consolidated=True, use_cftime=True)
cat = open_catalog("https://raw.githubusercontent.com/pangeo-data/pangeo-datastore/master/intake-catalogs/ocean/GFDL_CM2.6.yaml")
ds_ocean  = cat["GFDL_CM2_6_control_ocean_surface"].to_dask()
ds_atmos = xr.open_zarr('gs://cmip6/GFDL_CM2_6/control/atmos_daily.zarr', **kwargs)

regridder = xe.Regridder(ds_atmos, ds.surface_temp, 'bilinear')

my kernel crashes each time. I think the issue is memory useage, but this even crashes with the large server config. I found this and was wondering if this is a viable option to run this in the cloud. In the end we might want to precompute the weights here to make this operation easier to perform?

Thoughts?

rabernat commented 2 years ago

Good idea. But you will have to adapt that approach considerably for it to work in the cloud. There is no MPI available in this environment.

jbusecke commented 2 years ago

Good idea. But you will have to adapt that approach considerably for it to work in the cloud. There is no MPI available in this environment.

Sounds hard. Would it be easier to get higher memory servers in the deployment?

I am actually running into even more issues here which confuse me.

So taking the setup from above I am getting around the memory issue by subsetting the ocean dataset

lon_select = [-90, 0]
lat_select = [0, 45]
# ds_atmos = ds_atmos.sel(grid_xt=slice(lon_select[0], lon_select[1]), grid_yt=slice(lat_select[0], lat_select[1]))
ds = ds[['surface_temp']].sel(xt_ocean=slice(lon_select[0], lon_select[1]), yt_ocean=slice(lat_select[0], lat_select[1]))
ds

then creating a regridder object and adding the variables of interest to the main (ocean) dataset:

regridder = xe.Regridder(ds_atmos, ds.surface_temp, 'bilinear', periodic=False)

for var in ['wind', 't_ref', 'q_ref', 'ps']:
    ds[var] = regridder(ds_atmos[var])

Now I tried to look at a sample of the regridded data:

sample = ds_atmos.wind.isel(time=0)

If I load the data before regridding, I can easily plot the results:

regridder(sample.load()).plot()

image

But if I try to apply the regridding lazily (and load after the call to regridder) it never finishes.

regridder(sample).load().plot()

I thought xESMF would be able to be applied fully lazy? What could be going on here?

jbusecke commented 2 years ago

So the latter problem did not occur when running on gyre. Ill investigate more soon.

jbusecke commented 2 years ago

I am unsure what caused my initial issue here, but I havent had any problems with xesmf lately, so Ill close this here.