dt.to_zarr() fails with upscaled dimensions

[iskandari]

I am primarily using zarr to store 4D spatialtemporal data, and was initially inspired by the work of carbonplan. As I scale the temporal dimension of xarray DataTree objects, I end up generating very large graphs which crash dask when I attempt writing out a DataTree object with dt.to_zarr()

Following this 4D example, I get aDataTree object with 6 groups containing following dimensions and chunk sizes:

for i in range(0,6):
    print(dt[f"/{str(i)}"].dims)
    print(dt[f"/{str(i)}"]['climate'].data.chunks)

Frozen({'x': 128, 'y': 128, 'month': 12, 'band': 2})
((2,), (12,), (128,), (128,))
Frozen({'x': 256, 'y': 256, 'month': 12, 'band': 2})
((2,), (12,), (128, 128), (128, 128))
Frozen({'x': 512, 'y': 512, 'month': 12, 'band': 2})
((2,), (12,), (128, 128, 128, 128), (128, 128, 128, 128))
Frozen({'x': 1024, 'y': 1024, 'month': 12, 'band': 2})
((2,), (12,), (128, 128, 128, 128, 128, 128, 128, 128), (128, 128, 128, 128, 128, 128, 128, 128))
Frozen({'x': 2048, 'y': 2048, 'month': 12, 'band': 2})
((2,), (12,), (128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128), (128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128))
Frozen({'x': 4096, 'y': 4096, 'month': 12, 'band': 2})
((2,), (12,), (128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128), (128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128))

Using this structure, the largest dataset at level 6 is 1.5 GB , which I am able write to a Zarr pyramid without any issues:

dt.to_zarr(store_4d, consolidated=True, mode="w")

However, keeping everything else equal, if I scale the temporal dimension (month) by a factor for 4 - from 12 to 48 months - the dask graph size becomes too large (6GB):

%%time
# open and extract the input datasets
ds1_all = []
ds2_all = []
months = list(map(lambda d: d + 1, range(48))) #extend month range by a factor of 4
for i in months:
    path = f"{input_path}/wc2.1_2.5m_tavg_{(12 if i%12 == 0 else i%12):02g}.tif"
    ds = (
        xr.open_dataarray(path, engine="rasterio")
        .to_dataset(name="climate")
        .squeeze()
        .reset_coords(["band"], drop=True)
    )
    ds1_all.append(ds)
ds1 = xr.concat(ds1_all, pd.Index(months, name="month"))
for i in months:
    path = f"{input_path}/wc2.1_2.5m_prec_{(12 if i%12 == 0 else i%12):02g}.tif"
    ds = (
        xr.open_dataarray(path, engine="rasterio")
        .to_dataset(name="climate")
        .squeeze()
        .reset_coords(["band"], drop=True)
    )
    ds2_all.append(ds)


ds2 = xr.concat(ds2_all, pd.Index(months, name="month"))
ds2["climate"].values[ds2["climate"].values == ds2["climate"].values[0, 0, 0]] = ds1[
    "climate"
].values[0, 0, 0]
ds = xr.concat([ds1, ds2], pd.Index(["tavg", "prec"], name="band"))

# create the pyramid
dt = pyramid_reproject(
    ds, levels=LEVELS, extra_dim="band", other_chunks={'band': 2, 'month': 12}, clear_attrs=True
)

dt.to_zarr(store_4d, consolidated=True, mode="w")

UserWarning: Sending large graph of size 24.02 MiB.
UserWarning: Sending large graph of size 96.04 MiB.
UserWarning: Sending large graph of size 384.09 MiB.
UserWarning: Sending large graph of size 1.50 GiB.
UserWarning: Sending large graph of size 6.00 GiB.

FutureCancelledError: ('store-map-68286dca2cdf3de3e2a996a987b31cfe', 0, 3, 13, 20) cancelled for reason: scheduler-connection-lost.

At level 6, the 6 GB graph crashes dask.
I am running on a 12-core M2 macOS with 64 GiB RAM, and my dask configuration is:

client = Client(n_workers=4, threads_per_worker=2, memory_limit='15GB')

Some of the things I have tried:

• Starting the dask client with less workers/ higher memory limit (e.g. 1 worker with 64GB limit) and more workers/lower memory limit (e.g. 12 workers with 5GB limit), but the error is the same

• Setting higher array chunk sizes

from dask import config
config.set({'array.chunk-size': '10GiB'})

• Re-chunking the DataTree before writing to zarr

dt['/5'] = dt['/5'].chunk({'band': 2, 'month': 12, 'y': 1024, 'x': 1024})
dt['/5'].to_zarr(store_4d, consolidated=True, mode="w")

• Deferring the computation by writing a dask.delayed.Delayed object, but this does not seem to be implemented yet for DataTree objects

dt.to_zarr(store_4d, consolidated=True, compute=False, mode="w")
write_job = write_job.persist()
distributed.progress(write_job, notebook=False)

    187     raise NotImplementedError(
    188         "specifying a root group for the tree has not been implemented"
    189     )
    191 if not kwargs.get("compute", True):
--> 192     raise NotImplementedError("compute=False has not been implemented yet")
    194 if encoding is None:
    195     encoding = {}

NotImplementedError: compute=False has not been implemented yet

None of these worked though. Will it be possible in zarr 3.0 to write out DataTrees to zarr in parallel?

Are there any suggestions or xarray best practices (dask configuration, limiting graph size) for how to write out Zarr data pyramids without running into these errors for larger dimensions? Is it expected that the task graph size remains the same despite re-chunking? I would assume the specs of my machine should be sufficient to handle writing a 6GB Zarr pyramids at the scale listed in the above example, but I am not sure in which direction to go.

[iskandari]

Update: on the latest v3.0.0-beta.1, writing any DataTree object to zarr fails with

ValueError(
        "compressor cannot be used for arrays with version 3. Use bytes-to-bytes codecs instead."
    )

even though no compressor is specified