[WIP]: Force Zarr coordinate reads to be on the host

[TomAugspurger]

(just a POC. I'm not sure if this is a good idea or not)

zarr-python 3.x supports reading data to host (CPU) memory or device (GPU) memory. Because coordinates are small and really do need to be on the host (IIUC because putting them in an Index) then there's no benefit to reading them to device.

zarr-python includes a global config for whether to use host or device memory for reads, with zarr.config.enable_gpu(). But you can override that on a per-read basis by passing prototype to the getitem call.

This does that for arrays that are coordinates.

Here's a snippet:

import pathlib
import xarray as xr
import zarr


def main():
    p = pathlib.Path("test.zarr")
    if not p.exists():
        xr.tutorial.open_dataset("air_temperature").to_zarr("test.zarr", zarr_format=3)

    with zarr.config.enable_gpu():
        ds = xr.open_dataset("test.zarr", engine="zarr")
        print(ds)

        # This hits a TypeError in NumpyIndexingAdapter
        print(type(ds.air.data))


if __name__ == "__main__":
    main()

This still fails without #10078, but a little later on when we try to load the data. All the coordinate loading (including decode_times) seems to work fine. With this + #10078, the output is as expected:

<xarray.Dataset> Size: 31MB
Dimensions:  (time: 2920, lat: 25, lon: 53)
Coordinates:
  * time     (time) datetime64[ns] 23kB 2013-01-01 ... 2014-12-31T18:00:00
  * lat      (lat) float32 100B 75.0 72.5 70.0 67.5 65.0 ... 22.5 20.0 17.5 15.0
  * lon      (lon) float32 212B 200.0 202.5 205.0 207.5 ... 325.0 327.5 330.0
Data variables:
    air      (time, lat, lon) float64 31MB ...
Attributes:
    Conventions:  COARDS
    title:        4x daily NMC reanalysis (1948)
    description:  Data is from NMC initialized reanalysis\n(4x/day).  These a...
    platform:     Model
    references:   http://www.esrl.noaa.gov/psd/data/gridded/data.ncep.reanaly...
<class 'cupy.ndarray'>  # <-------- this is the important bit!

cc @dcherian

[negin513]

+1 . Coordinates should stay on the host and "data" should be streamed to device. I would not think metada (including coordinates) to be read directly on GPUs ever and in general they are so small that there will not be any additional benefit in reading them to device.

[TomAugspurger]

I added a new keyword coords_buffer_prototype to the open_zarr API. By default, that's None which means use zarr.core.buffer.cpu.buffer_prototype (i.e. force it to be on the host). If you really need it to be on the GPU or some other device, then you can provide the buffer prototype to use for the host here, and it should be forwarded all the way through to the zarr.Array.getitem call.

I'll try adding a test here.

[dcherian]

Thanks for opening this, now I see what you mean.

Let me push an alternative PR. We only need this because our decoding machinery does handle duck arrays yet. So at the cost of a device->host transfer of a small array, we can fix it in the decoding pipeline.