Add more unit tests for GPU buffer

[maxrjones]

This PR adds some tests for the GPU buffer prototype to unlock #2738.

[maxrjones]

@TomAugspurger would you expect the sharding codec to work with the gpu buffer prototype? If so, I could debug the failure more but don't want to spend time setting up an environment with cuda unnecessarily

[TomAugspurger]

I haven't used or looked at the sharding code, so I don't know offhand sorry.

If setting up a dev end with cupy is a hassle I can take a look sometime.

Just looking at the traceback, maybe one thing to check: try setting with zarr.config.enable_gpu() in the tests? Maybe we're grabbing a buffer type from the default config somewhere where we should be basing it off another Buffer class?

[jakirkham]

would you expect the sharding codec to work with the gpu buffer prototype?

cc @madsbk @akshaysubr (in case either of you have thoughts on this question)

[akshaysubr]

Its been a while since I looked at the sharding codec, but it was not expected to work with the GPU buffer prototype. Might be worth another closer look but it seemed like some things in the sharding codec need to still be generalized to work with arbitrary buffer prototypes.

[d-v-b]

the sharding codec is unique because it does IO, and thus interacts with a store, which might be the root of the complication here. (I would argue it is not really a codec in the conventional sense, and is rather a special implementation of an array).

[TomAugspurger]

Looking into this a bit today. For some reason the get_ndbuffer_class is returning a CPU buffer in sharding.py:

-> await get_pipeline_class()
(Pdb) ll
650         async def _encode_shard_index(self, index: _ShardIndex) -> Buffer:
651             index_bytes = next(
652                 iter(
653  ->                 await get_pipeline_class()
654                     .from_codecs(self.index_codecs)
655                     .encode(
656                         [
657                             (
658                                 get_ndbuffer_class().from_numpy_array(index.offsets_and_lengths),
659                                 self._get_index_chunk_spec(index.chunks_per_shard),
660                             )
661                         ],
662                     )
663                 )
664             )
665             assert index_bytes is not None
666             assert isinstance(index_bytes, Buffer)
667             return index_bytes
(Pdb) pp get_ndbuffer_class()
<class 'zarr.core.buffer.cpu.NDBuffer'>

despite us setting it to gpu.NDBuffer in the test setup. I'll push a fix later.

[TomAugspurger]

188e501 seems to fix that test.

I notice a bunch of other uses of numpy_buffer_prototype() in that file, which AFAICT is using host memory by design. We'll might need to audit this a bit further to see what all can / should be using default_buffer_prototype() instead.

[TomAugspurger]

@maxrjones is this in a good state now from your side?

[maxrjones]

@maxrjones is this in a good state now from your side?

Thanks for looking into this @TomAugspurger and pushing a fix! Yes, seems good to me. I agree with looking into the other numpy_buffer_prototypes() calls separately.