Add Kvikio backend entrypoint

[dcherian]

TODO:

• Avoid calling np.asarray on lazy indexing classes pydata/xarray#6874
• Remove np.asarray in formatting.py pydata/xarray#8100
• Generalize explicit_indexing_adapter pydata/xarray#8408
• add docs - https://cupy-xarray--10.org.readthedocs.build/generated/cupy_xarray.kvikio.KvikioBackendEntrypoint.html
• add tests

This PR registers a "kvikio" backend that allows us to use kvikio to read from network into GPU memory directly using "GDS".

To get this to work we need

1. [REVIEW] Support of alternative array classes zarr-developers/zarr-python#934
2. Avoid calling np.asarray on lazy indexing classes pydata/xarray#6874
3. kvikio.zarr working (https://github.com/rapidsai/kvikio)

This PR subclasses the existing ZarrStore and overrides necessary methods. most of the code is actually copied over from xarray/backends/zarr.py

For a short demo see https://github.com/dcherian/cupy-xarray/blob/kvikio-entrypoint/docs/source/kvikio.ipynb

[andersy005]

Xarray + Zarr + GPUs!! So exciting...

[dcherian]

So excited I couldn't wait to hack it together :P

[weiji14]

Cool, great work @dcherian, it's like Christmas came early! I've tried to test this branch but am encountering some cuFile issues. Will just post what I've done so far in case anyone wants to reproduce this branch:

# May need to install nvidia-gds first
# https://docs.nvidia.com/cuda/cuda-installation-guide-linux/index.html#ubuntu-installation-common
sudo apt install nvidia-gds

git clone https://github.com/dcherian/cupy-xarray.git
cd cupy-xarray

mamba create --name cupy-xarray python=3.9 cupy=11.0 rapidsai-nightly::kvikio=22.10 jupyterlab=3.4.5 pooch=1.6.0 netcdf4=1.6.0 watermark=2.3.1
mamba activate cupy-xarray
python -m ipykernel install --user --name cupy-xarray

# https://github.com/pydata/xarray/pull/6874
pip install git+https://github.com/dcherian/xarray.git@kvikio
# https://github.com/zarr-developers/zarr-python/pull/934
# pip install git+https://github.com/madsbk/zarr-python.git@cupy_support
pip install zarr==2.13.0a2
# https://github.com/xarray-contrib/cupy-xarray/pull/10
git switch kvikio-entrypoint
pip install --editable=.

# Start jupyter lab
jupyter lab --no-browser
# Then open the docs/kvikio.ipynb notebook

With that, I got an error on the ds = xr.open_dataset(store, engine="kvikio", consolidated=False) cell, full traceback as follows:

---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
Input In [4], in <cell line: 2>()
      1 # Consolidated must be False
----> 2 ds = xr.open_dataset(store, engine="kvikio", consolidated=False)
      3 print(ds.air._variable._data)
      4 ds

File ~/mambaforge/envs/cupy-xarray/lib/python3.9/site-packages/xarray/backends/api.py:531, in open_dataset(filename_or_obj, engine, chunks, cache, decode_cf, mask_and_scale, decode_times, decode_timedelta, use_cftime, concat_characters, decode_coords, drop_variables, inline_array, backend_kwargs, **kwargs)
    519 decoders = _resolve_decoders_kwargs(
    520     decode_cf,
    521     open_backend_dataset_parameters=backend.open_dataset_parameters,
   (...)
    527     decode_coords=decode_coords,
    528 )
    530 overwrite_encoded_chunks = kwargs.pop("overwrite_encoded_chunks", None)
--> 531 backend_ds = backend.open_dataset(
    532     filename_or_obj,
    533     drop_variables=drop_variables,
    534     **decoders,
    535     **kwargs,
    536 )
    537 ds = _dataset_from_backend_dataset(
    538     backend_ds,
    539     filename_or_obj,
   (...)
    547     **kwargs,
    548 )
    549 return ds

File ~/Documents/github/cupy-xarray/cupy_xarray/kvikio.py:190, in KvikioBackendEntrypoint.open_dataset(self, filename_or_obj, mask_and_scale, decode_times, concat_characters, decode_coords, drop_variables, use_cftime, decode_timedelta, group, mode, synchronizer, consolidated, chunk_store, storage_options, stacklevel)
    188 store_entrypoint = StoreBackendEntrypoint()
    189 with close_on_error(store):
--> 190     ds = store_entrypoint.open_dataset(
    191         store,
    192         mask_and_scale=mask_and_scale,
    193         decode_times=decode_times,
    194         concat_characters=concat_characters,
    195         decode_coords=decode_coords,
    196         drop_variables=drop_variables,
    197         use_cftime=use_cftime,
    198         decode_timedelta=decode_timedelta,
    199     )
    200 return ds

File ~/mambaforge/envs/cupy-xarray/lib/python3.9/site-packages/xarray/backends/store.py:26, in StoreBackendEntrypoint.open_dataset(self, store, mask_and_scale, decode_times, concat_characters, decode_coords, drop_variables, use_cftime, decode_timedelta)
     14 def open_dataset(
     15     self,
     16     store,
   (...)
     24     decode_timedelta=None,
     25 ):
---> 26     vars, attrs = store.load()
     27     encoding = store.get_encoding()
     29     vars, attrs, coord_names = conventions.decode_cf_variables(
     30         vars,
     31         attrs,
   (...)
     38         decode_timedelta=decode_timedelta,
     39     )

File ~/mambaforge/envs/cupy-xarray/lib/python3.9/site-packages/xarray/backends/common.py:125, in AbstractDataStore.load(self)
    103 def load(self):
    104     """
    105     This loads the variables and attributes simultaneously.
    106     A centralized loading function makes it easier to create
   (...)
    122     are requested, so care should be taken to make sure its fast.
    123     """
    124     variables = FrozenDict(
--> 125         (_decode_variable_name(k), v) for k, v in self.get_variables().items()
    126     )
    127     attributes = FrozenDict(self.get_attrs())
    128     return variables, attributes

File ~/mambaforge/envs/cupy-xarray/lib/python3.9/site-packages/xarray/backends/zarr.py:461, in ZarrStore.get_variables(self)
    460 def get_variables(self):
--> 461     return FrozenDict(
    462         (k, self.open_store_variable(k, v)) for k, v in self.zarr_group.arrays()
    463     )

File ~/mambaforge/envs/cupy-xarray/lib/python3.9/site-packages/xarray/core/utils.py:474, in FrozenDict(*args, **kwargs)
    473 def FrozenDict(*args, **kwargs) -> Frozen:
--> 474     return Frozen(dict(*args, **kwargs))

File ~/mambaforge/envs/cupy-xarray/lib/python3.9/site-packages/xarray/backends/zarr.py:462, in <genexpr>(.0)
    460 def get_variables(self):
    461     return FrozenDict(
--> 462         (k, self.open_store_variable(k, v)) for k, v in self.zarr_group.arrays()
    463     )

File ~/Documents/github/cupy-xarray/cupy_xarray/kvikio.py:130, in GDSZarrStore.open_store_variable(self, name, zarr_array)
    128 else:
    129     array_wrapper = CupyZarrArrayWrapper
--> 130 data = indexing.LazilyIndexedArray(array_wrapper(name, self))
    132 attributes = dict(attributes)
    133 encoding = {
    134     "chunks": zarr_array.chunks,
    135     "preferred_chunks": dict(zip(dimensions, zarr_array.chunks)),
    136     "compressor": zarr_array.compressor,
    137     "filters": zarr_array.filters,
    138 }

File ~/mambaforge/envs/cupy-xarray/lib/python3.9/site-packages/xarray/backends/zarr.py:64, in ZarrArrayWrapper.__init__(self, variable_name, datastore)
     61 self.datastore = datastore
     62 self.variable_name = variable_name
---> 64 array = self.get_array()
     65 self.shape = array.shape
     67 dtype = array.dtype

File ~/Documents/github/cupy-xarray/cupy_xarray/kvikio.py:34, in EagerCupyZarrArrayWrapper.get_array(self)
     33 def get_array(self):
---> 34     return np.asarray(self)

File ~/Documents/github/cupy-xarray/cupy_xarray/kvikio.py:31, in EagerCupyZarrArrayWrapper.__array__(self)
     30 def __array__(self):
---> 31     return self.datastore.zarr_group[self.variable_name][:].get()

File ~/mambaforge/envs/cupy-xarray/lib/python3.9/site-packages/zarr/core.py:807, in Array.__getitem__(self, selection)
    805     result = self.vindex[selection]
    806 else:
--> 807     result = self.get_basic_selection(pure_selection, fields=fields)
    808 return result

File ~/mambaforge/envs/cupy-xarray/lib/python3.9/site-packages/zarr/core.py:933, in Array.get_basic_selection(self, selection, out, fields)
    930     return self._get_basic_selection_zd(selection=selection, out=out,
    931                                         fields=fields)
    932 else:
--> 933     return self._get_basic_selection_nd(selection=selection, out=out,
    934                                         fields=fields)

File ~/mambaforge/envs/cupy-xarray/lib/python3.9/site-packages/zarr/core.py:976, in Array._get_basic_selection_nd(self, selection, out, fields)
    970 def _get_basic_selection_nd(self, selection, out=None, fields=None):
    971     # implementation of basic selection for array with at least one dimension
    972 
    973     # setup indexer
    974     indexer = BasicIndexer(selection, self)
--> 976     return self._get_selection(indexer=indexer, out=out, fields=fields)

File ~/mambaforge/envs/cupy-xarray/lib/python3.9/site-packages/zarr/core.py:1267, in Array._get_selection(self, indexer, out, fields)
   1261 if not hasattr(self.chunk_store, "getitems") or \
   1262    any(map(lambda x: x == 0, self.shape)):
   1263     # sequentially get one key at a time from storage
   1264     for chunk_coords, chunk_selection, out_selection in indexer:
   1265 
   1266         # load chunk selection into output array
-> 1267         self._chunk_getitem(chunk_coords, chunk_selection, out, out_selection,
   1268                             drop_axes=indexer.drop_axes, fields=fields)
   1269 else:
   1270     # allow storage to get multiple items at once
   1271     lchunk_coords, lchunk_selection, lout_selection = zip(*indexer)

File ~/mambaforge/envs/cupy-xarray/lib/python3.9/site-packages/zarr/core.py:1966, in Array._chunk_getitem(self, chunk_coords, chunk_selection, out, out_selection, drop_axes, fields)
   1962 ckey = self._chunk_key(chunk_coords)
   1964 try:
   1965     # obtain compressed data for chunk
-> 1966     cdata = self.chunk_store[ckey]
   1968 except KeyError:
   1969     # chunk not initialized
   1970     if self._fill_value is not None:

File ~/mambaforge/envs/cupy-xarray/lib/python3.9/site-packages/zarr/storage.py:1066, in DirectoryStore.__getitem__(self, key)
   1064 filepath = os.path.join(self.path, key)
   1065 if os.path.isfile(filepath):
-> 1066     return self._fromfile(filepath)
   1067 else:
   1068     raise KeyError(key)

File ~/mambaforge/envs/cupy-xarray/lib/python3.9/site-packages/kvikio/zarr.py:43, in GDSStore._fromfile(self, fn)
     41 else:
     42     nbytes = os.path.getsize(fn)
---> 43     with kvikio.CuFile(fn, "r") as f:
     44         ret = cupy.empty(nbytes, dtype="u1")
     45         read = f.read(ret)

File ~/mambaforge/envs/cupy-xarray/lib/python3.9/site-packages/kvikio/cufile.py:67, in CuFile.__init__(self, file, flags)
     49 def __init__(self, file: Union[pathlib.Path, str], flags: str = "r"):
     50     """Open and register file for GDS IO operations
     51 
     52     CuFile opens the file twice and maintains two file descriptors.
   (...)
     65         "+" -> "open for updating (reading and writing)"
     66     """
---> 67     self._handle = libkvikio.CuFile(file, flags)

File libkvikio.pyx:103, in kvikio._lib.libkvikio.CuFile.__init__()

RuntimeError: cuFile error at: /workspace/.conda-bld/work/cpp/include/kvikio/file_handle.hpp:172: internal error

Interestingly, nvidia-smi is showing some GPU memory being allocated already 0 N/A N/A 11852 C ...vs/cupy-xarray/bin/python 187MiB, but I think there's some CUDA driver issue I might need to sort out first. Will do a bit more troubleshooting 😄

Edit: Looking at https://docs.nvidia.com/gpudirect-storage/release-notes/index.html#known-limitations and https://docs.nvidia.com/gpudirect-storage/release-notes/index.html#support-matrix, it seems that gpudirect-storage is only supported for certain filesystems like ext4, but I have btrfs on my laptop 😅. Let me try it on my university HPC 🙂

Edit2: Nevermind, I just remembered my root folder is ext4 so I just put the Zarr file in /tmp. Now it works!!

[joshmoore]

As in xarray-contrib/xbatcher#87 (comment)

Now available in zarr-python 2.13.0a2 for testing.

[jakirkham]

Should add 2.13 final is out as well

[dcherian]

As an update, we need more upstream work in Xarray: pydata/xarray#8100

and some work to dask to get that to work (see the notebook)

@negin513 This should be OK to try on derecho now. Dask will not work, but everything else should.

[weiji14]

Maybe we can refactor this to use kvikio.zarr.open_cupy_array once kvikio=23.10 is out? There's support for nvCOMP-based LZ4 compression now (that's compatible with Zarr's CPU-based LZ4 compressor), xref rapidsai/kvikio#267.

[dcherian]

PR welcome!

[weiji14]

PR welcome!

Debating on whether to start from scratch in a completely new branch, or rebase off of this one 😄

P.S. I'm starting some work over at https://github.com/weiji14/foss4g2023oceania for a conference talk on 18 Oct, hoping to get the kvikIO engine running on an ERA5 dataset. Somehow, I could only get things on this branch to work up to kvikIO==23.06, it seems like RAPIDS AI 23.08 moved to CUDA 12 and I've been getting some errors like RuntimeError: Unable to open file: Too many open files. I'll try to push some code in the next 2 weeks, and should be able to nudge this forward a little bit 😉

[dcherian]

Just build on top of this branch in a new PR.

The optimization i was mentioning to save self.datastore.zarr_group[self.variable_name] as self._array. Otherwise we keep pinging the store unnecessarily

[weiji14]

Ok, I got a bit too ambitious trying to work in support for LZ4 compression (via nvCOMP), and hit into some issues (see rapidsai/kvikio#297). There's some stable releases of RAPIDS AI kvikIO 23.10.00 and xarray 2023.10.0 now which should have some nice enhancements for this PR. I'll try to squeeze out some time to work on it.

[dcherian]

Cool blogpost @weiji14 ! Where are you running your ERA5 benchmark?

This kvikio engine can be optimized quite a bit given the work we (Earthmover) have recently done on the Xarray Zarr backend. There's a number of unnecessary metadata reads that can be removed.

[weiji14]

Cool blogpost @weiji14 ! Where are you running your ERA5 benchmark?

Thanks! I'm running the benchmark locally on my laptop (with an RTX A2000 8GB GPU), see also zarr-developers/zarr-benchmark#14, and I've got a few more extra technical details mentioned at my own site https://weiji14.github.io/blog/when-cloud-native-geospatial-meets-gpu-native-machine-learning 😉

This kvikio engine can be optimized quite a bit given the work we (Earthmover) have recently done on the Xarray Zarr backend. There's a number of unnecessary metadata reads that can be removed.

Yes, I saw https://earthmover.io/blog/cloud-native-dataloader, and think a lot of the pieces are coming together. We really should wrap up this PR at some point, I'd be happy to give this another review if you'd like to push more changes on this branch.

[weiji14]

Currently stuck on this line, and trying to figure out what to change in relation to pydata/xarray#6874. My guess is that in an ideal world where we've solved all the numpy/cupy duck typing issues, we could just use the recently modified ZarrArrayWrapper (pydata/xarray#8472) like this:

Suggested change

	data = indexing.LazilyIndexedArray(array_wrapper(zarr_array))
	data = indexing.LazilyIndexedArray(ZarrArrayWrapper(zarr_array))

However, that leads to this error:

____________________________________________________________________________________ test_lazy_load[True] ____________________________________________________________________________________

consolidated = True, store = PosixPath('/tmp/pytest-of-weiji/pytest-32/test_lazy_load_True_0/kvikio.zarr')

    @pytest.mark.parametrize("consolidated", [True, False])
    def test_lazy_load(consolidated, store):
>       with xr.open_dataset(store, engine="kvikio", consolidated=consolidated) as ds:

cupy_xarray/tests/test_kvikio.py:37: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
../../../mambaforge/envs/cupy-xarray-doc/lib/python3.10/site-packages/xarray/backends/api.py:571: in open_dataset
    backend_ds = backend.open_dataset(
cupy_xarray/kvikio.py:241: in open_dataset
    ds = store_entrypoint.open_dataset(
../../../mambaforge/envs/cupy-xarray-doc/lib/python3.10/site-packages/xarray/backends/store.py:58: in open_dataset
    ds = Dataset(vars, attrs=attrs)
../../../mambaforge/envs/cupy-xarray-doc/lib/python3.10/site-packages/xarray/core/dataset.py:711: in __init__
    variables, coord_names, dims, indexes, _ = merge_data_and_coords(
../../../mambaforge/envs/cupy-xarray-doc/lib/python3.10/site-packages/xarray/core/dataset.py:425: in merge_data_and_coords
    return merge_core(
../../../mambaforge/envs/cupy-xarray-doc/lib/python3.10/site-packages/xarray/core/merge.py:699: in merge_core
    collected = collect_variables_and_indexes(aligned, indexes=indexes)
../../../mambaforge/envs/cupy-xarray-doc/lib/python3.10/site-packages/xarray/core/merge.py:362: in collect_variables_and_indexes
    idx, idx_vars = create_default_index_implicit(variable)
../../../mambaforge/envs/cupy-xarray-doc/lib/python3.10/site-packages/xarray/core/indexes.py:1404: in create_default_index_implicit
    index = PandasIndex.from_variables(dim_var, options={})
../../../mambaforge/envs/cupy-xarray-doc/lib/python3.10/site-packages/xarray/core/indexes.py:654: in from_variables
    obj = cls(data, dim, coord_dtype=var.dtype)
../../../mambaforge/envs/cupy-xarray-doc/lib/python3.10/site-packages/xarray/core/indexes.py:589: in __init__
    index = safe_cast_to_index(array)
../../../mambaforge/envs/cupy-xarray-doc/lib/python3.10/site-packages/xarray/core/indexes.py:469: in safe_cast_to_index
    index = pd.Index(np.asarray(array), **kwargs)
../../../mambaforge/envs/cupy-xarray-doc/lib/python3.10/site-packages/xarray/core/indexing.py:509: in __array__
    return np.asarray(self.get_duck_array(), dtype=dtype)
../../../mambaforge/envs/cupy-xarray-doc/lib/python3.10/site-packages/xarray/backends/common.py:181: in get_duck_array
    return self[key]  # type: ignore [index]
../../../mambaforge/envs/cupy-xarray-doc/lib/python3.10/site-packages/xarray/backends/zarr.py:104: in __getitem__
    return indexing.explicit_indexing_adapter(
../../../mambaforge/envs/cupy-xarray-doc/lib/python3.10/site-packages/xarray/core/indexing.py:1017: in explicit_indexing_adapter
    indexable = NumpyIndexingAdapter(result)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

self = <[AttributeError("'NumpyIndexingAdapter' object has no attribute 'array'") raised in repr()] NumpyIndexingAdapter object at 0x7f71beb14a00>
array = array([-5, -4, -3, -2, -1,  0,  1,  2,  3,  4])

    def __init__(self, array):
        # In NumpyIndexingAdapter we only allow to store bare np.ndarray
        if not isinstance(array, np.ndarray):
>           raise TypeError(
                "NumpyIndexingAdapter only wraps np.ndarray. "
                f"Trying to wrap {type(array)}"
            )
E           TypeError: NumpyIndexingAdapter only wraps np.ndarray. Trying to wrap <class 'cupy.ndarray'>

../../../mambaforge/envs/cupy-xarray-doc/lib/python3.10/site-packages/xarray/core/indexing.py:1493: TypeError
================================================================================== short test summary info ===================================================================================
FAILED cupy_xarray/tests/test_kvikio.py::test_lazy_load[True] - TypeError: NumpyIndexingAdapter only wraps np.ndarray. Trying to wrap <class 'cupy.ndarray'>
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!! stopping after 1 failures !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

Looked into this a bit and got lost in all the xarray class inheritance logic and mixins... @dcherian, could you point out what's the recommended route forward? Is there something we still need to upstream into xarray? I'm testing this on xarray=2024.6.0 btw.

Output of xr.show_versions() for reference:

INSTALLED VERSIONS
------------------
commit: None
python: 3.10.14 | packaged by conda-forge | (main, Mar 20 2024, 12:45:18) [GCC 12.3.0]
python-bits: 64
OS: Linux
OS-release: 6.7.12+bpo-amd64
machine: x86_64
processor: 
byteorder: little
LC_ALL: None
LANG: en_NZ.UTF-8
LOCALE: ('en_NZ', 'UTF-8')
libhdf5: 1.14.3
libnetcdf: 4.9.2

xarray: 2024.6.0
pandas: 2.2.2
numpy: 1.26.4
scipy: None
netCDF4: 1.7.1
pydap: None
h5netcdf: None
h5py: None
zarr: 2.18.2
cftime: 1.6.4
nc_time_axis: None
iris: None
bottleneck: None
dask: 2024.6.2
distributed: None
matplotlib: None
cartopy: None
seaborn: None
numbagg: None
fsspec: 2024.6.0
cupy: 13.2.0
pint: None
sparse: None
flox: None
numpy_groupies: None
setuptools: 70.0.0
pip: 24.0
conda: None
pytest: 8.2.2
mypy: None
IPython: 8.25.0
sphinx: 7.3.7