Lazily caching computations in xarray/dask?

[SohumB]

I tried searching for this, but I'm new to the community, and I'm not even sure what term we would use for this behaviour. If I've missed something obvious please let me know!

Say I have a computation specified on an xarray dataset backed by a zarr store.

ds = xarray.open_zarr('s3://...')
rs = ds.tasmax.resample(time = "1Y").max(dim="time")

At this point, nothing is computed, right? And as I understand it, the usual pattern for using this rs object is to either compute the whole object and persist it, or to use as the input for future computations that may eventually apply a filter, and that filter can get propagated up to ensure that only the appropriate chunks get opened from the s3 store.

Can I do something in between? Can I request that rs be "in principle" backed by a different remote zarr store, and when some child computation of rs applies a filter and forces a computation, if that second store doesn't contain the necessary data, to compute it in this fashion and also store it?

ds = xarray.open_zarr('s3://...')
rs = ds.tasmax.resample(time="1Y").max(dim="time").backed_by('s3://...')
... 
rs.isel(...).compute() # this triggers a cache check, and if it fails, a compute and store

[dcherian]

With the single machine scheduler dask can optionally cache: https://docs.dask.org/en/stable/caching.html. This won't persist across sessions.

You could also try https://github.com/ncar-xdev/xpersist

[SohumB]

Thanks for the pointer! Unfortunately I don't think either of those fit — the goal here is to persist a partially computed xarray across sessions without fully computing it, basically. I don't think that's what xpersist does; though I'm not sure — it seems based on xr.backends.api.to_zarr, which I think will fully force any xarray you hand it?

[dcherian]

Hmm maybe graphchain? I'd think of this as more of a dask problem and ask on that repo instead of here.

[SohumB]

Ah, that could help, certainly! And that makes sense, I'll go ask the question to them.

Thanks!