Skip to content

Commit 2f9cf22

Browse files
akshaysubrd-v-b
akshaysubr
and
d-v-b
authored
Initial GPU support (#1967)
* Initial implementation of a GPU version of Buffer and NDBuffer * Adding cupy as an optional dependency * Adding GPU prototype test * Adding GPU memory store implementation * Addressing comments * Making GpuMemoryStore tests conditional on cupy being available * Adding test checking that existing host memory codecs use the gpu_buffer_prototype appropriately * Reducing code and docs duplication * Formatting * Fixing silent rebase conflicts * Reducing code duplication in GpuMemoryStore * Refactoring to an abstract Buffer class and concrete CPU and GPU implementations of those * Templating store tests on Buffer type * Changing imports to prevent circular dependencies * Fixing unsafe calls to Buffer abstract methods in metadata.py and group.py * Preventing calls to abstract classmethods of Buffer and NDBuffer * Fixing some more unsafe usage of Buffer abstract class * Initial testing with cirun based GPU CI * Reverting to basic ubuntu machine image on GCP * Switching to cuda image from the docker registry * Revert "Switching to cuda image from the docker registry" This reverts commit d473a3d. * Revert "Reverting to basic ubuntu machine image on GCP" This reverts commit e5cfd2f. * Revert "Initial testing with cirun based GPU CI" This reverts commit ff40d3c. * Adding pytest mark for GPU tests * Updating GPU memory store test with gpu mark * Adding GPU workflow that only runs GPU tests * Formatting * Fixing mypy errors in buffer code * Fixing errors in test_buffer.py * Fixing errors in test_buffer.py * Fixing store test errors * Fixing stateful store test * Fixing config test * Fixing group tests * Fixing indexing tests * Manually installing cupy in the GPU workflow * Ablating GPU test matrix and adding gpu optional dependencies to the hatch dependency list * Adding some more logging to debug GPU test failures * Adding GA step to install the CUDA toolkit * Adding a separate gputest hatch environment to simplify GPU testing * Fixing error in cuda-toolkit step * Downgrading to CUDA 12.4.1 in cuda-toolkit GA * Trying manual install of the CUDA toolkit * Updating environment variables with CUDA installation * Removing PATH env and setting it only through GITHUB_PATH * Fixing issue from merge conflict --------- Co-authored-by: Davis Bennett <[email protected]>
1 parent 6ea9997 commit 2f9cf22

24 files changed

+853
-115
lines changed

‎.github/workflows/gpu_test.yml‎

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,66 @@
1+
# This workflow will install Python dependencies, run tests and lint with a variety of Python versions
2+
# For more information see: https://help.github.com/actions/language-and-framework-guides/using-python-with-github-actions
3+
4+
name: GPU Test V3
5+
6+
on:
7+
push:
8+
branches: [ v3 ]
9+
pull_request:
10+
branches: [ v3 ]
11+
workflow_dispatch:
12+
13+
env:
14+
LD_LIBRARY_PATH: /usr/local/cuda/extras/CUPTI/lib64:/usr/local/cuda/lib64
15+
16+
concurrency:
17+
group: ${{ github.workflow }}-${{ github.ref }}
18+
cancel-in-progress: true
19+
20+
jobs:
21+
test:
22+
name: py=${{ matrix.python-version }}, np=${{ matrix.numpy-version }}, deps=${{ matrix.dependency-set }}
23+
24+
runs-on: gpu-runner
25+
strategy:
26+
matrix:
27+
python-version: ['3.11']
28+
numpy-version: ['2.0']
29+
dependency-set: ["minimal"]
30+
31+
steps:
32+
- uses: actions/checkout@v4
33+
# - name: cuda-toolkit
34+
# uses: Jimver/[email protected]
35+
# id: cuda-toolkit
36+
# with:
37+
# cuda: '12.4.1'
38+
- name: Set up CUDA
39+
run: |
40+
wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/cuda-keyring_1.1-1_all.deb
41+
sudo dpkg -i cuda-keyring_1.1-1_all.deb
42+
sudo apt-get update
43+
sudo apt-get -y install cuda-toolkit-12-6
44+
echo "/usr/local/cuda/bin" >> $GITHUB_PATH
45+
- name: GPU check
46+
run: |
47+
nvidia-smi
48+
echo $PATH
49+
echo $LD_LIBRARY_PATH
50+
nvcc -V
51+
- name: Set up Python
52+
uses: actions/setup-python@v5
53+
with:
54+
python-version: ${{ matrix.python-version }}
55+
cache: 'pip'
56+
- name: Install Hatch and CuPy
57+
run: |
58+
python -m pip install --upgrade pip
59+
pip install hatch
60+
- name: Set Up Hatch Env
61+
run: |
62+
hatch env create gputest.py${{ matrix.python-version }}-${{ matrix.numpy-version }}-${{ matrix.dependency-set }}
63+
hatch env run -e gputest.py${{ matrix.python-version }}-${{ matrix.numpy-version }}-${{ matrix.dependency-set }} list-env
64+
- name: Run Tests
65+
run: |
66+
hatch env run --env gputest.py${{ matrix.python-version }}-${{ matrix.numpy-version }}-${{ matrix.dependency-set }} run-coverage

‎pyproject.toml‎

Lines changed: 34 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -74,6 +74,9 @@ jupyter = [
7474
'ipytree>=0.2.2',
7575
'ipywidgets>=8.0.0',
7676
]
77+
gpu = [
78+
"cupy-cuda12x",
79+
]
7780
docs = [
7881
'sphinx',
7982
'sphinx-autobuild>=2021.3.14',
@@ -120,7 +123,7 @@ build.hooks.vcs.version-file = "src/zarr/_version.py"
120123
[tool.hatch.envs.test]
121124
dependencies = [
122125
"numpy~={matrix:numpy}",
123-
"universal_pathlib"
126+
"universal_pathlib",
124127
]
125128
features = ["test", "extra"]
126129

@@ -134,8 +137,34 @@ python = ["3.10", "3.11", "3.12"]
134137
numpy = ["1.24", "1.26", "2.0"]
135138
features = ["optional"]
136139

140+
[[tool.hatch.envs.test.matrix]]
141+
python = ["3.10", "3.11", "3.12"]
142+
numpy = ["1.24", "1.26", "2.0"]
143+
features = ["gpu"]
144+
137145
[tool.hatch.envs.test.scripts]
138146
run-coverage = "pytest --cov-config=pyproject.toml --cov=pkg --cov=tests"
147+
run-coverage-gpu = "pip install cupy-cuda12x && pytest -m gpu --cov-config=pyproject.toml --cov=pkg --cov=tests"
148+
run = "run-coverage --no-cov"
149+
run-verbose = "run-coverage --verbose"
150+
run-mypy = "mypy src"
151+
run-hypothesis = "pytest --hypothesis-profile ci tests/v3/test_properties.py tests/v3/test_store/test_stateful*"
152+
list-env = "pip list"
153+
154+
[tool.hatch.envs.gputest]
155+
dependencies = [
156+
"numpy~={matrix:numpy}",
157+
"universal_pathlib",
158+
]
159+
features = ["test", "extra", "gpu"]
160+
161+
[[tool.hatch.envs.gputest.matrix]]
162+
python = ["3.10", "3.11", "3.12"]
163+
numpy = ["1.24", "1.26", "2.0"]
164+
version = ["minimal"]
165+
166+
[tool.hatch.envs.gputest.scripts]
167+
run-coverage = "pytest -m gpu --cov-config=pyproject.toml --cov=pkg --cov=tests"
139168
run = "run-coverage --no-cov"
140169
run-verbose = "run-coverage --verbose"
141170
run-mypy = "mypy src"
@@ -223,4 +252,8 @@ filterwarnings = [
223252
"error:::zarr.*",
224253
"ignore:PY_SSIZE_T_CLEAN will be required.*:DeprecationWarning",
225254
"ignore:The loop argument is deprecated since Python 3.8.*:DeprecationWarning",
255+
"ignore:Creating a zarr.buffer.gpu.*:UserWarning",
256+
]
257+
markers = [
258+
"gpu: mark a test as requiring CuPy and GPU"
226259
]

‎src/zarr/codecs/blosc.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,8 @@
1010

1111
from zarr.abc.codec import BytesBytesCodec
1212
from zarr.core.array_spec import ArraySpec
13-
from zarr.core.buffer import Buffer, as_numpy_array_wrapper
13+
from zarr.core.buffer import Buffer
14+
from zarr.core.buffer.cpu import as_numpy_array_wrapper
1415
from zarr.core.common import JSON, parse_enum, parse_named_configuration, to_thread
1516
from zarr.registry import register_codec
1617

‎src/zarr/codecs/gzip.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,8 @@
77

88
from zarr.abc.codec import BytesBytesCodec
99
from zarr.core.array_spec import ArraySpec
10-
from zarr.core.buffer import Buffer, as_numpy_array_wrapper
10+
from zarr.core.buffer import Buffer
11+
from zarr.core.buffer.cpu import as_numpy_array_wrapper
1112
from zarr.core.common import JSON, parse_named_configuration, to_thread
1213
from zarr.registry import register_codec
1314

‎src/zarr/codecs/zstd.py‎

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

1010
from zarr.abc.codec import BytesBytesCodec
1111
from zarr.core.array_spec import ArraySpec
12-
from zarr.core.buffer import Buffer, as_numpy_array_wrapper
12+
from zarr.core.buffer import Buffer
13+
from zarr.core.buffer.cpu import as_numpy_array_wrapper
1314
from zarr.core.common import JSON, parse_named_configuration, to_thread
1415
from zarr.registry import register_codec
1516

‎src/zarr/core/array.py‎

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -512,15 +512,24 @@ async def _set_selection(
512512

513513
# check value shape
514514
if np.isscalar(value):
515-
value = np.asanyarray(value, dtype=self.metadata.dtype)
515+
array_like = prototype.buffer.create_zero_length().as_array_like()
516+
if isinstance(array_like, np._typing._SupportsArrayFunc):
517+
# TODO: need to handle array types that don't support __array_function__
518+
# like PyTorch and JAX
519+
array_like_ = cast(np._typing._SupportsArrayFunc, array_like)
520+
value = np.asanyarray(value, dtype=self.metadata.dtype, like=array_like_)
516521
else:
517522
if not hasattr(value, "shape"):
518523
value = np.asarray(value, self.metadata.dtype)
519524
# assert (
520525
# value.shape == indexer.shape
521526
# ), f"shape of value doesn't match indexer shape. Expected {indexer.shape}, got {value.shape}"
522527
if not hasattr(value, "dtype") or value.dtype.name != self.metadata.dtype.name:
523-
value = np.array(value, dtype=self.metadata.dtype, order="A")
528+
if hasattr(value, "astype"):
529+
# Handle things that are already NDArrayLike more efficiently
530+
value = value.astype(dtype=self.metadata.dtype, order="A")
531+
else:
532+
value = np.array(value, dtype=self.metadata.dtype, order="A")
524533
value = cast(NDArrayLike, value)
525534
# We accept any ndarray like object from the user and convert it
526535
# to a NDBuffer (or subclass). From this point onwards, we only pass

‎src/zarr/core/buffer/__init__.py‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,19 @@
1+
from zarr.core.buffer.core import (
2+
ArrayLike,
3+
Buffer,
4+
BufferPrototype,
5+
NDArrayLike,
6+
NDBuffer,
7+
default_buffer_prototype,
8+
)
9+
from zarr.core.buffer.cpu import numpy_buffer_prototype
10+
11+
__all__ = [
12+
"ArrayLike",
13+
"Buffer",
14+
"NDArrayLike",
15+
"NDBuffer",
16+
"BufferPrototype",
17+
"default_buffer_prototype",
18+
"numpy_buffer_prototype",
19+
]

0 commit comments

Comments
 (0)