Usage of RAM will explode when performing calculation on files processed by kerchunk

Hi all,

Thank you for developing this fantastic library that allows for direct chunking of netcdf files! However, I noticed some unusual behavior when conducting calculations on the file after it has been processed by kerchunk (verision: 0.0.9). Any comments and suggestions are well appreciated!

Background

I show an example on Cheyenne that aims to extract the first layer of temperature (SST) from the CESM ocean model. The model time slice files have been concated into a time series file ('/glade/campaign/univ/umia0022/time_series/e.e12.E20TRC5CN.f09_g16_EMOM_001/tseries/ocn/monthly/e.e12.E20TRC5CN.f09_g16_EMOM_001.EMOM.h0.mon.TEMP.192001-200512.nc').

A JSON file was created using kerchunk (/glade/scratch/che43/share/CESM1-EMOM-jsons/e.e12.E20TRC5CN.f09_g16_EMOM_001.EMOM.h0.mon.TEMP.192001-200512.nc.f09_g16_EMOM_001.json)

The comparison between accessing the netcdf file directly and using fsspec is presented below.

Import library and initialize cluster

import warnings
warnings.filterwarnings("ignore")
import xarray as xr
import fsspec
import json
from tqdm import tqdm
from glob import glob
from ncar_jobqueue import NCARCluster
import dask
import ujson



num_jobs = 20
cluster = NCARCluster(cores=num_jobs,
                      processes=1, 
                      memory='25GB', 
                      project='UOSU0006', 
                      queue='casper',
                      walltime="6:00:00")
cluster.scale(jobs=num_jobs)

from distributed import Client
from distributed.utils import format_bytes

client = Client(cluster)
cluster

open the data

ds is based on fssepc, and ds1 is directly opened.

var = 'TEMP'
fnames = glob('/glade/scratch/che43/share/CESM1-EMOM-jsons/e.e12.E20TRC5CN.f09_g16_EMOM_001.EMOM.h0.mon.TEMP.192001-200512.nc.f09_g16_EMOM_001.json')
json_files = [ fname for fname in fnames if '.'+var+'.' in fname]
json_files = sorted(json_files)


m_list = []
for js in tqdm(json_files):
    with open(js) as f:
        m_list.append(fsspec.get_mapper("reference://", 
            fo=ujson.load(f), remote_protocol="file",
            remote_options={"anon": True}))


ds = xr.open_mfdataset(m_list[0:1], engine='zarr', 
                        combine='nested', coords='minimal', 
                      chunks={'time':-1, 'nlat':192, 'nlon':160, 'z_t':1},)

ds1 = xr.open_mfdataset(['/glade/campaign/univ/umia0022/time_series/e.e12.E20TRC5CN.f09_g16_EMOM_001/tseries/ocn/monthly/e.e12.E20TRC5CN.f09_g16_EMOM_001.EMOM.h0.mon.TEMP.192001-200512.nc'],
                       chunks={'time':-1, 'nlat':192, 'nlon':160, 'z_t':1},)

They look like the same
ds

ds1

extract the first layer

%%time
ds1.TEMP.isel(z_t=0).compute()

ds1 successfully completes the task. Each worker only uses memory less than 1GB.

%%time
# the RAM usage will explode for workers and cannot proceed.
ds.TEMP.isel(z_t=0).compute()

Memory explodes.

Start by reproducing the example with kerchunk 0.0.9, the fsspec reference mapper, xarray open_mfdataset, and the Dask-backed .compute() calls, comparing the kerchunk dataset with the direct NetCDF dataset. Identify why selecting z_t=0 causes worker memory to grow, then verify that the same calculation completes without the reported memory explosion.