Skip to content

Commit 33a9184

Browse files
committed
handle warnings in gpu tests
1 parent 8afdce9 commit 33a9184

2 files changed

Lines changed: 19 additions & 13 deletions

File tree

tests/test_buffer.py

Lines changed: 15 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
from zarr.codecs.gzip import GzipCodec
1414
from zarr.codecs.transpose import TransposeCodec
1515
from zarr.codecs.zstd import ZstdCodec
16+
from zarr.errors import ZarrUserWarning
1617
from zarr.storage import MemoryStore, StorePath
1718
from zarr.testing.buffer import (
1819
NDBufferUsingTestNDArrayLike,
@@ -138,12 +139,13 @@ async def test_codecs_use_of_gpu_prototype() -> None:
138139
filters=[TransposeCodec(order=(1, 0))],
139140
)
140141
expect[:] = cp.arange(100).reshape(10, 10)
141-
142-
await a.setitem(
143-
selection=(slice(0, 10), slice(0, 10)),
144-
value=expect[:],
145-
prototype=gpu.buffer_prototype,
146-
)
142+
msg = "Creating a zarr.buffer.gpu.Buffer with an array that does not support the __cuda_array_interface__ for zero-copy transfers, falling back to slow copy based path"
143+
with pytest.warns(ZarrUserWarning, match=msg):
144+
await a.setitem(
145+
selection=(slice(0, 10), slice(0, 10)),
146+
value=expect[:],
147+
prototype=gpu.buffer_prototype,
148+
)
147149
got = await a.getitem(selection=(slice(0, 10), slice(0, 10)), prototype=gpu.buffer_prototype)
148150
assert isinstance(got, cp.ndarray)
149151
assert cp.array_equal(expect, got)
@@ -164,12 +166,13 @@ async def test_sharding_use_of_gpu_prototype() -> None:
164166
fill_value=0,
165167
)
166168
expect[:] = cp.arange(100).reshape(10, 10)
167-
168-
await a.setitem(
169-
selection=(slice(0, 10), slice(0, 10)),
170-
value=expect[:],
171-
prototype=gpu.buffer_prototype,
172-
)
169+
msg = "Creating a zarr.buffer.gpu.Buffer with an array that does not support the __cuda_array_interface__ for zero-copy transfers, falling back to slow copy based path"
170+
with pytest.warns(ZarrUserWarning, match=msg):
171+
await a.setitem(
172+
selection=(slice(0, 10), slice(0, 10)),
173+
value=expect[:],
174+
prototype=gpu.buffer_prototype,
175+
)
173176
got = await a.getitem(
174177
selection=(slice(0, 10), slice(0, 10)), prototype=gpu.buffer_prototype
175178
)

tests/test_store/test_memory.py

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99

1010
import zarr
1111
from zarr.core.buffer import Buffer, cpu, gpu
12+
from zarr.errors import ZarrUserWarning
1213
from zarr.storage import GpuMemoryStore, MemoryStore
1314
from zarr.testing.store import StoreTests
1415
from zarr.testing.utils import gpu_test
@@ -130,6 +131,8 @@ def test_from_dict(self) -> None:
130131
"a": gpu.Buffer.from_bytes(b"aaaa"),
131132
"b": cpu.Buffer.from_bytes(b"bbbb"),
132133
}
133-
result = GpuMemoryStore.from_dict(d)
134+
msg = "Creating a zarr.buffer.gpu.Buffer with an array that does not support the __cuda_array_interface__ for zero-copy transfers, falling back to slow copy based path"
135+
with pytest.warns(ZarrUserWarning, match=msg):
136+
result = GpuMemoryStore.from_dict(d)
134137
for v in result._store_dict.values():
135138
assert type(v) is gpu.Buffer

0 commit comments

Comments
 (0)