Multi-GPU setup and python/zarr api

[merrysailor]

Hi,

First, the project is excellent, very helpful even in my setup which uses the posix fallback.

Second, I want to highlight that having a zarr decompress directly to the desired device_id is a bit baroque, as the device_ordinal used during compression is used by default unless overwritten explicitly:

zarr.open_array(
    kvikio.zarr.GDSStore(
        path=f,
        decompressor_config_overwrite=kvikio.zarr.Snappy(device_ordinal=1).get_config()
    ),
    mode='r',
    meta_array=cupy.empty(()),
)[...]

Btw, I confirm that with the above the cupy array ends up in the correct device. But wanted to double check that indeed the above works as expected all the way, i.e., the compressed data move and decompression indeed happen on the specified device?

Third, the above fails when multiple python threads are involved. Not familiar with CUDA, but appears that streams are thread specific? When trying to use the above in a multi-threaded setup (have a threaded data loader):

File "libnvcomp.pyx", line 246, in kvikio._lib.libnvcomp._SnappyManager.__cinit__
RuntimeError: nvCOMP error: Stream 0 is not associated with device 1

Thank you.

[madsbk]

Second, I want to highlight that having a zarr decompress directly to the desired device_id is a bit baroque, as the device_ordinal used during compression is used by default unless overwritten explicitly.

Yeah, we haven't really supported this use case :/ The new nvCOMP should automatically choose the correct device based on the stream however our nvCOMP bindings doesn't support stream. I am removing the options in #650.

But we are working on adding KvikIO/nvCOMP support directly in Zarr v3 (first step: zarr-developers/zarr-python#2863), which hopefully makes this much easier!

[merrysailor]

Thank you, I will keep an eye on this.