Support CuPy-backed arrays in DaskManager#11383
Conversation
The "meta" argument passed to dask.array.from_array should not be hardcoded to just `numpy.ndarray`, but allow for `cupy.ndarray` too.
Not sure how to type-hint np | cp | ??, so just use Any for output of get_array_namespace.
| # lazily loaded backend array classes should use NumPy array operations. | ||
| kwargs["meta"] = np.ndarray | ||
| # lazily loaded backend array classes should use NumPy or CuPy array operations. | ||
| xp = get_array_namespace(data.get_duck_array()) |
There was a problem hiding this comment.
we probably need to add some API allow getting the underlying array type / library without actually fetching data. Something like a data.get_array_namespace() or data.get_meta()? Not sure how easy it would be to implement that, though.
(I think this is what causes the tests to fail)
There was a problem hiding this comment.
Ideally we could just call xp = data.__array_namespace__() following the Array API spec - https://data-apis.org/array-api/2025.12/API_specification/generated/array_api.array.__array_namespace__.html, and it would propagate through all the subclassed layers to get the underlying array namespace (numpy or cupy). I thought this would work by putting it into the NDArrayMixin (b77cc57), but that breaks a lot of the lazy repls...
Might need to think this through a bit more. Wondering if there needs to be a cached .__cached_array_namespace__ attribute of some sort to work with the lazy objects...
There was a problem hiding this comment.
i went down this road a while ago. the problem is that we need our lazy arrays coerced to an in-memory type at some point. If they advertise __array_namespace__ they can treated as an in-memory type (deep within dask), and nothing works. Did you also run in to this problem?
There was a problem hiding this comment.
I think I've went pretty deep, but not deep into dask internals (yet?). If it ends up needing changes in dask, I'm just gonna push on findind a way to remove dask entirely - #9038 (comment).
Right now I think I've solved the repl issues by changing some of the logic in formatting.py. Still need to work my way through some other logic that have been basing their logic around duck arrays / __array_namespace__ but shouldn't...
There was a problem hiding this comment.
Actually:
[ ] Ensure calling
.compute()on an xarray.Dataset backed by CuPy arrays don't get coerced to NumPy arrays.
I did realize that even with the cupy meta fixes here, the Dask-backed arrays still get loaded into NumPy instead. So maybe that's what you're referring to? That there's no way to have a fully CuPy-only array pipeline without changing some internals of Dask?
There was a problem hiding this comment.
Is there an example I can run without an nvidia gpu?
There was a problem hiding this comment.
Unfortunately no, I only know how to reproduce this on this cupy-xarray PR - https://github.com/xarray-contrib/cupy-xarray/pull/81/changes#diff-b866be1141ec4295c6a2ef9b8effab4a58d6dacd262918d0db4f56d6d7927818:
ds: xr.Dataset = xr.open_mfdataset(
paths=[
"https://github.com/developmentseed/titiler/raw/1.2.0/src/titiler/mosaic/tests/fixtures/B01.tif",
"https://github.com/developmentseed/titiler/raw/1.2.0/src/titiler/mosaic/tests/fixtures/B09.tif",
],
engine="cog3pio",
concat_dim="band",
combine="nested",
device_id=None,
)
ds.raster.load() # inplace load, requires https://github.com/pydata/xarray/pull/11381
# assert isinstance(
# ds.raster, cp.ndarray # TODO wait for https://github.com/pydata/xarray/pull/11383 ?
# )But I should figure out a good duck array test somehow, and will let you know if it can be reproduced without a GPU.
Centralize retrieving of the __array_namespace__ through several subclassed layers, to avoid having to go through `.get_duck_array()`. Need to put `from xarray.compat.array_api_compat import get_array_namespace` import within the method to avoid circular import. Also type-hinted output of `get_array_namespace` as ModuleType following numpy/numpy#20719.
To fix repr AssertionError mismatches on: - TestVariable::test_repr_lazy_data - test_repr_pandas_multi_index - test_repr_pandas_range_index - test_display_nbytes - test_repr_file_collapsed - test_coordinate_transform_variable_repr by preventing specific xarray internal array types from going through the is_duck_array repl path.
| return short_array_repr(array) | ||
| elif is_duck_array(internal_data): | ||
| elif not isinstance( | ||
| internal_data, (LazilyIndexedArray, MemoryCachedArray, IndexingAdapter) |
There was a problem hiding this comment.
| internal_data, (LazilyIndexedArray, MemoryCachedArray, IndexingAdapter) | |
| internal_data, ExplicitlyIndexed |
is usually the way to do it
There was a problem hiding this comment.
Excellent hint, but actually, not needed anymore. I talked to @keewis on the Pangeo call just now, and figured out the right-level of abstraction is to put the __array_namespace__ method in ImplicitToExplicitIndexingAdapter instead of NDArrayMixin (commit 7132f73). So we won't need to change this repr stuff anymore.
There was a problem hiding this comment.
Oh wait, commit 7132f73 might not be correct either 😅 I'll actually need to use your suggestion here, and really should write a proper test for this first...
Move `__array_namespace__` method from NDArrayMixin to just ImplicitToExplicitIndexingAdapter.
|
CI is failing with Still need to make sure calling |
3 participants
Description
The default Dask ChunkManagerEntrypoint appears to be hardcoded to return NumPy arrays by default, even if the underlying arrays are CuPy arrays
TODO:
dask.array.from_arraythat was hardcoded tonp.ndarray(since 2019 in Ensure explicitly indexed arrays are preserved #3027), allowingcupy.ndarrayinstead when detected.compute()on an xarray.Dataset backed by CuPy arrays don't get coerced to NumPy arrays.Probably needs #11381 to be merged first. Part of resolving xarray-contrib/cupy-xarray#81 (comment).
Checklist
whats-new.rstapi.rstAI Disclosure