Add CuPy support to ArrayStore via Array API (#579)

## Description

<!-- Provide a brief description of the PR's purpose here. -->

This PR supports CuPy via the Array API. At first I thought this might
be challenging since CuPy has a different model of devices, i.e., it
depends on context managers rather than passing in devices. However,
array-api-compat seems to have solved these issues and made CuPy
compatible in terms of devices. Here are some references that I looked
at:

- data-apis/array-api-compat#293
-
https://github.com/data-apis/array-api-compat/blob/main/array_api_compat/cupy/_aliases.py
-
https://docs.cupy.dev/en/stable/reference/array_api_functions.html#cupy.array_api.ones

## TODO

<!-- Notable points that this PR has either accomplished or will
accomplish. -->

- [x] Add cupy backend to conftest
- [x] Add instructions for installing cupy
- [x] Test ArrayStore with CuPy on GPU
- [x] Update ArrayStore to consider CuPy in pandas return type
- [x] Add cupy tests to CI -- I installed CuPy but the CuPy tests won't
actually run because we don't have a GPU runner set up -- following
scipy
(https://github.com/scipy/scipy/blob/main/.github/workflows/gpu-ci.yml),
I also set the CI to only run on ubuntu, since installing cupy on mac
and windows seems to be problematic

## Status

- [x] I have read the guidelines in

[CONTRIBUTING.md](https://github.com/icaros-usc/pyribs/blob/master/CONTRIBUTING.md)
- [x] I have formatted my code using `yapf`
- [x] I have tested my code by running `pytest`
- [x] I have linted my code with `pylint`
- [x] I have added a one-line description of my change to the changelog
in
      `HISTORY.md`
- [x] This PR is ready to go

Author: btjanaka

.github/workflows/testing.yml

@@ -65,7 +65,7 @@ jobs:
     strategy:
       max-parallel: 12 # All in parallel.
       matrix:
-        os: [ubuntu-latest, macos-latest, windows-latest]
+        os: [ubuntu-latest]
         python-version: ["3.12"]
     runs-on: ${{ matrix.os }}
     steps:
@@ -77,6 +77,8 @@ jobs:
         run: pip install .[dev]
       - name: Install torch
         run: pip install torch
+      - name: Install cupy
+        run: pip install cupy-cuda12x
       - name: Test core
         run: >
           pytest tests/archives tests/emitters tests/schedulers

HISTORY.md

@@ -6,7 +6,7 @@
 
 #### API
 
-- Support array backends via Python array API Standard ({pr}`573`, {pr}`571`)
+- Support array backends via Python array API Standard ({issue}`570`)
 - **Backwards-incompatible:** Remove raw_dict methods from ArrayStore
   ({pr}`575`)
 

README.md

@@ -124,10 +124,11 @@ python -c "import ribs; print(ribs.__version__)"
 You should see a version number in the output.
 
 **Experimental:** Pyribs is experimenting with adding support for running QD
-algorithms in PyTorch via the
+algorithms in PyTorch and CuPy via the
 [Python array API standard](https://data-apis.org/array-api/latest/). To enable
-this functionality, [install PyTorch](https://pytorch.org), such as with
-`pip install torch`.
+this functionality, install [PyTorch](https://pytorch.org), such as with
+`pip install torch`, and [CuPy](https://cupy.dev), such as with
+`pip install cupy-cuda12x`.
 
 ## Usage
 

ribs/archives/_array_store.py

@@ -4,7 +4,13 @@
 from enum import IntEnum
 from functools import cached_property
 
-from array_api_compat import is_numpy_array, is_numpy_namespace, is_torch_array
+from array_api_compat import (is_cupy_array, is_numpy_array, is_numpy_namespace,
+                              is_torch_array)
+
+try:
+    from array_api_compat import cupy as cp
+except ImportError:
+    pass
 
 from ribs._utils import arr_readonly, xp_namespace
 from ribs.archives._archive_data_frame import ArchiveDataFrame
@@ -292,10 +298,12 @@ def _convert_to_numpy(arr):
             return arr
         elif is_torch_array(arr):
             return arr.cpu().detach().numpy()
+        elif is_cupy_array(arr):
+            return cp.asnumpy(arr)
         else:
             raise NotImplementedError(
                 "The pandas return type is currently only supported "
-                "with numpy and torch arrays.")
+                "with NumPy, PyTorch, and CuPy arrays.")
 
     def retrieve(self, indices, fields=None, return_type="dict"):
         """Collects data at the given indices.

tests/conftest.py

@@ -24,6 +24,21 @@
 except ImportError:
     pass
 
+try:
+    import cupy as cp
+    from cupy_backends.cuda.api.runtime import CUDARuntimeError
+
+    cp.empty(0)  # Triggers CUDARuntimeError if there is no GPU available.
+
+    xp_available_backends.append(
+        pytest.param((cp, cp.cuda.Device(0)), id="cupy-gpu"))
+except ImportError:
+    # CuPy not installed.
+    pass
+except CUDARuntimeError:
+    # GPU not available.
+    pass
+
 
 @pytest.fixture(params=xp_available_backends)
 def xp_and_device(request):