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.