Debug cubed + NLDAS virtual store on the VEDA Hub

[maxrjones]

Opening a ticket for tracking - I tried to use coiled with the NLDAS virtual store on the VEDA Hub and it fails silently. I think this may be related to a Hub issue (https://developmentseed.slack.com/archives/C04R9B69T5W/p1764968671305229).

Here's a MVCe (notice that "Cubed success" is never printed):

# /// script
# requires-python = ">=3.10"
# dependencies = [
#     "icechunk",
#     "xarray",
#     "zarr",
#     "cubed",
#     "cubed-xarray",
#     "dask",
# ]
# ///
print("starting")
import icechunk
import xarray as xr
import warnings

warnings.filterwarnings(
    "ignore",
    message="Numcodecs codecs are not in the Zarr version 3 specification*",
    category=UserWarning
)

# Open icechunk repository
storage = icechunk.s3_storage(
    bucket='nasa-waterinsight',
    prefix=f"virtual-zarr-store/NLDAS-3-icechunk",
    region="us-west-2",
    anonymous=True,
)

chunk_url = "s3://nasa-waterinsight/NLDAS3/forcing/daily/"
virtual_credentials = icechunk.containers_credentials({
    chunk_url: icechunk.s3_anonymous_credentials()
})

repo = icechunk.Repository.open(
    storage=storage,
    authorize_virtual_chunk_access=virtual_credentials,
)

session = repo.readonly_session('main')
ds = xr.open_zarr(session.store, consolidated=False, zarr_format=3, chunks={}, chunked_array_type="dask")
tair = ds["Tair"].sel(time="2001-01-02", lat=slice(10, 15), lon=slice(-60, -55))
data = tair.load()
print("Dask success")

ds = xr.open_zarr(session.store, consolidated=False, zarr_format=3, chunks={}, chunked_array_type="cubed")
tair = ds["Tair"].sel(time="2001-01-02", lat=slice(10, 15), lon=slice(-60, -55))
data = tair.load()
print("Cubed success")

[jsignell]

I ran this on the hub and the kernel died because it had run out of memory. I think there are a few things going on here:

1. Cubed is requiring much more memory than dask to do the same operation
2. Cubed is supposed to raise an error ahead of time rather than crash the kernel

For 1 there is likely some optimization that dask implements and cubed does not. This might have to do with how slicing is pushed down to the array layer.

For 2 there are many things that could be going on, but my suspicion is that cubed isn't able to get a good sense of the memory available on jupyterhub.

[jsignell]

I noticed that there is a way of specifying the memory limit explicitly to cubed so I tried adding that in:

spec = cubed.Spec(allowed_mem="2GB")

 ds = xr.open_zarr(
    session.store, consolidated=False, zarr_format=3, chunks={},
    chunked_array_type="cubed", from_array_kwargs=dict(spec=spec)
)

But it didn't seem to help at all and I ended up bailing before my instance ran out of memory.

[maxrjones]

For 2 there are many things that could be going on, but my suspicion is that cubed isn't able to get a good sense of the memory available on jupyterhub.

This doesn't surprise me. Many system utilities report the node memory availability rather than the pod memory availability (e.g., top, htop).

spec = cubed.Spec(allowed_mem="2GB")

do you remember what size instance you were using for this test?

[jsignell]

spec = cubed.Spec(allowed_mem="2GB")

do you remember what size instance you were using for this test?

Probably the 7.4GB option, but I just ran it again with the 14.8GB option and that ran out of memory as well.

[jsignell]

Wellll can't claim I did anything, but I just remembered about this ticket and tried the reproducer again and everything worked smoothly. The memory didn't skyrocket, it finished in a reasonable amount of time. I imagine the issue was resolved upstream in cubed (maybe by cubed-dev/cubed#858)