Merge branch 'ig/typesize_for_blosc' of github.com:ilan-gold/zarr-python into ig/typesize_for_blosc

Author: ilan-gold

changes/2978.bugfix.rst

@@ -0,0 +1 @@
+Fixed sharding with GPU buffers.

pyproject.toml

@@ -83,7 +83,7 @@ test = [
 remote_tests = [
     'zarr[remote]',
     "botocore",
-    "s3fs",
+    "s3fs>=2023.10.0",
     "moto[s3,server]",
     "requests",
 ]
@@ -104,7 +104,7 @@ docs = [
     # Optional dependencies to run examples
     'numcodecs[msgpack]',
     'rich',
-    's3fs',
+    's3fs>=2023.10.0',
     'astroid<4'
 ]
 
@@ -348,27 +348,27 @@ python_version = "3.11"
 ignore_missing_imports = true
 namespace_packages = false
 
-
 strict = true
 warn_unreachable = true
-
 enable_error_code = ["ignore-without-code", "redundant-expr", "truthy-bool"]
 
+
 [[tool.mypy.overrides]]
 module = [
-    "zarr.v2.*",
+    "tests.package_with_entrypoint.*",
+    "tests.test_codecs.test_transpose",
+    "tests.test_config"
 ]
-ignore_errors = true
+strict = false
 
+# TODO: Move the next modules up to the strict = false section
+# and fix the errors
 [[tool.mypy.overrides]]
 module = [
     "zarr.testing.stateful",  # lots of hypothesis decorator errors
-    "tests.package_with_entrypoint.*",
     "tests.test_codecs.test_codecs",
-    "tests.test_codecs.test_transpose",
     "tests.test_metadata.*",
     "tests.test_store.*",
-    "tests.test_config",
     "tests.test_group",
     "tests.test_indexing",
     "tests.test_properties",

src/zarr/codecs/sharding.py

@@ -683,7 +683,7 @@ def _get_index_chunk_spec(self, chunks_per_shard: ChunkCoords) -> ArraySpec:
             config=ArrayConfig(
                 order="C", write_empty_chunks=False
             ),  # Note: this is hard-coded for simplicity -- it is not surfaced into user code,
-            prototype=numpy_buffer_prototype(),
+            prototype=default_buffer_prototype(),
         )
 
     def _get_chunk_spec(self, shard_spec: ArraySpec) -> ArraySpec:

src/zarr/testing/utils.py

@@ -44,7 +44,7 @@ def has_cupy() -> bool:
 # Decorator for GPU tests
 def gpu_test(func: T_Callable) -> T_Callable:
     return cast(
-        T_Callable,
+        "T_Callable",
         pytest.mark.gpu(
             pytest.mark.skipif(not has_cupy(), reason="CuPy not installed or no GPU available")(
                 func

tests/package_with_entrypoint/__init__.py

@@ -1,13 +1,14 @@
 from collections.abc import Iterable
+from typing import Any
 
-from numpy import ndarray
+import numpy as np
+import numpy.typing as npt
 
 import zarr.core.buffer
-from zarr.abc.codec import ArrayBytesCodec, CodecInput, CodecOutput, CodecPipeline
+from zarr.abc.codec import ArrayBytesCodec, CodecInput, CodecPipeline
 from zarr.codecs import BytesCodec
 from zarr.core.array_spec import ArraySpec
 from zarr.core.buffer import Buffer, NDBuffer
-from zarr.core.common import BytesLike
 
 
 class TestEntrypointCodec(ArrayBytesCodec):
@@ -16,14 +17,14 @@ class TestEntrypointCodec(ArrayBytesCodec):
     async def encode(
         self,
         chunks_and_specs: Iterable[tuple[CodecInput | None, ArraySpec]],
-    ) -> Iterable[CodecOutput | None]:
-        pass
+    ) -> Iterable[Buffer | None]:
+        return [None]
 
     async def decode(
         self,
         chunks_and_specs: Iterable[tuple[CodecInput | None, ArraySpec]],
-    ) -> ndarray:
-        pass
+    ) -> npt.NDArray[Any]:
+        return np.array(1)
 
     def compute_encoded_size(self, input_byte_length: int, chunk_spec: ArraySpec) -> int:
         return input_byte_length
@@ -35,13 +36,13 @@ def __init__(self, batch_size: int = 1) -> None:
 
     async def encode(
         self, chunks_and_specs: Iterable[tuple[CodecInput | None, ArraySpec]]
-    ) -> BytesLike:
-        pass
+    ) -> Iterable[Buffer | None]:
+        return [None]
 
     async def decode(
         self, chunks_and_specs: Iterable[tuple[CodecInput | None, ArraySpec]]
-    ) -> ndarray:
-        pass
+    ) -> Iterable[NDBuffer | None]:
+        return np.array(1)
 
 
 class TestEntrypointBuffer(Buffer):

tests/test_buffer.py

@@ -148,6 +148,34 @@ async def test_codecs_use_of_gpu_prototype() -> None:
     assert cp.array_equal(expect, got)
 
 
+@gpu_test
+@pytest.mark.asyncio
+async def test_sharding_use_of_gpu_prototype() -> None:
+    with zarr.config.enable_gpu():
+        expect = cp.zeros((10, 10), dtype="uint16", order="F")
+
+        a = await zarr.api.asynchronous.create_array(
+            StorePath(MemoryStore()) / "test_codecs_use_of_gpu_prototype",
+            shape=expect.shape,
+            chunks=(5, 5),
+            shards=(10, 10),
+            dtype=expect.dtype,
+            fill_value=0,
+        )
+        expect[:] = cp.arange(100).reshape(10, 10)
+
+        await a.setitem(
+            selection=(slice(0, 10), slice(0, 10)),
+            value=expect[:],
+            prototype=gpu.buffer_prototype,
+        )
+        got = await a.getitem(
+            selection=(slice(0, 10), slice(0, 10)), prototype=gpu.buffer_prototype
+        )
+        assert isinstance(got, cp.ndarray)
+        assert cp.array_equal(expect, got)
+
+
 def test_numpy_buffer_prototype() -> None:
     buffer = cpu.buffer_prototype.buffer.create_zero_length()
     ndbuffer = cpu.buffer_prototype.nd_buffer.create(shape=(1, 2), dtype=np.dtype("int64"))
@@ -157,6 +185,16 @@ def test_numpy_buffer_prototype() -> None:
         ndbuffer.as_scalar()
 
 
+@gpu_test
+def test_gpu_buffer_prototype() -> None:
+    buffer = gpu.buffer_prototype.buffer.create_zero_length()
+    ndbuffer = gpu.buffer_prototype.nd_buffer.create(shape=(1, 2), dtype=cp.dtype("int64"))
+    assert isinstance(buffer.as_array_like(), cp.ndarray)
+    assert isinstance(ndbuffer.as_ndarray_like(), cp.ndarray)
+    with pytest.raises(ValueError, match="Buffer does not contain a single scalar value"):
+        ndbuffer.as_scalar()
+
+
 # TODO: the same test for other buffer classes
 def test_cpu_buffer_as_scalar() -> None:
     buf = cpu.buffer_prototype.nd_buffer.create(shape=(), dtype="int64")

tests/test_codecs/test_transpose.py

@@ -48,6 +48,7 @@ async def test_transpose(
     read_data = await _AsyncArrayProxy(a)[:, :].get()
     assert np.array_equal(data, read_data)
 
+    assert isinstance(read_data, np.ndarray)
     if runtime_read_order == "F":
         assert read_data.flags["F_CONTIGUOUS"]
         assert not read_data.flags["C_CONTIGUOUS"]
@@ -90,5 +91,5 @@ def test_transpose_invalid(
                 dtype=data.dtype,
                 fill_value=0,
                 chunk_key_encoding={"name": "v2", "separator": "."},
-                filters=[TransposeCodec(order=order)],
+                filters=[TransposeCodec(order=order)],  # type: ignore[arg-type]
             )

tests/test_config.py

@@ -10,7 +10,7 @@
 import zarr
 import zarr.api
 from zarr import zeros
-from zarr.abc.codec import CodecInput, CodecOutput, CodecPipeline
+from zarr.abc.codec import CodecPipeline
 from zarr.abc.store import ByteSetter, Store
 from zarr.codecs import (
     BloscCodec,
@@ -21,6 +21,7 @@
 )
 from zarr.core.array_spec import ArraySpec
 from zarr.core.buffer import NDBuffer
+from zarr.core.buffer.core import Buffer
 from zarr.core.codec_pipeline import BatchedCodecPipeline
 from zarr.core.config import BadConfigError, config
 from zarr.core.indexing import SelectorTuple
@@ -144,7 +145,7 @@ def test_config_codec_pipeline_class(store: Store) -> None:
     class MockCodecPipeline(BatchedCodecPipeline):
         async def write(
             self,
-            batch_info: Iterable[tuple[ByteSetter, ArraySpec, SelectorTuple, SelectorTuple]],
+            batch_info: Iterable[tuple[ByteSetter, ArraySpec, SelectorTuple, SelectorTuple, bool]],
             value: NDBuffer,
             drop_axes: tuple[int, ...] = (),
         ) -> None:
@@ -174,7 +175,7 @@ async def write(
     class MockEnvCodecPipeline(CodecPipeline):
         pass
 
-    register_pipeline(MockEnvCodecPipeline)
+    register_pipeline(MockEnvCodecPipeline)  # type: ignore[type-abstract]
 
     with mock.patch.dict(
         os.environ, {"ZARR_CODEC_PIPELINE__PATH": fully_qualified_name(MockEnvCodecPipeline)}
@@ -191,10 +192,9 @@ def test_config_codec_implementation(store: Store) -> None:
     _mock = Mock()
 
     class MockBloscCodec(BloscCodec):
-        async def _encode_single(
-            self, chunk_data: CodecInput, chunk_spec: ArraySpec
-        ) -> CodecOutput | None:
+        async def _encode_single(self, chunk_bytes: Buffer, chunk_spec: ArraySpec) -> Buffer | None:
             _mock.call()
+            return None
 
     register_codec("blosc", MockBloscCodec)
     with config.set({"codecs.blosc": fully_qualified_name(MockBloscCodec)}):
@@ -245,7 +245,7 @@ def test_config_buffer_implementation() -> None:
     # has default value
     assert fully_qualified_name(get_buffer_class()) == config.defaults[0]["buffer"]
 
-    arr = zeros(shape=(100), store=StoreExpectingTestBuffer())
+    arr = zeros(shape=(100,), store=StoreExpectingTestBuffer())
 
     # AssertionError of StoreExpectingTestBuffer when not using my buffer
     with pytest.raises(AssertionError):