Repository: dask/dask
Pull request #11309 (2024-08-14 00:55:29)

Upgrade gpuCI and fix Dask Array failures with "cupy" backend
Related: https://github.com/dask/dask/issues/11312

- [ ] Closes #xxxx
- [ ] Tests added / passed
- [ ] Passes `pre-commit run --all-files`


github-actions[bot]: ## Unit Test Results
_See [test report](https://dask.github.io/dask/test_report.html) for an extended history of previous test failures. This is useful for diagnosing flaky tests._

     13 files  ±0       13 suites  ±0   2h 58m 28s :stopwatch: - 20m 33s
 13 149 tests ±0   12 079 :white_check_mark: ±0   1 070 :zzz: ±0  0 :x: ±0 
137 022 runs  ±0  117 931 :white_check_mark: ±0  19 091 :zzz: ±0  0 :x: ±0 

Results for commit 38628957. ± Comparison against base commit 67b2852d.

[test-results]:data:application/gzip;base64,H4sIAK4HvGYC/03Myw7CIBCF4VdpWLsYoBTGlzHcmhB7MRRWxneXWqRd/t9kzpuMYfIbuXeU3zqy5ZDOcjnqFNZlb5CgCpVrOu60x38/tmztjgzkBZ/hVT+bjTpMxU7wMa6xSszLMS2BsdptmkrktGGdRsBml+lfX5ftOs8hlSBcDUyhkGLQBnsDnAnj0RnwzAn03HOtnJOCfL7r3yjlGgEAAA==

rjzamora: This PR should resolve gpuCI failures once it is combined with https://github.com/dask/dask/pull/11242

Patch:
diff --git a/continuous_integration/gpuci/axis.yaml b/continuous_integration/gpuci/axis.yaml
index 25b4c78ef..7b30d5065 100644
--- a/continuous_integration/gpuci/axis.yaml
+++ b/continuous_integration/gpuci/axis.yaml
@@ -9,6 +9,6 @@ LINUX_VER:
 - ubuntu20.04
 
 RAPIDS_VER:
-- "24.08"
+- "24.10"
 
 excludes:
diff --git a/dask/array/_shuffle.py b/dask/array/_shuffle.py
index 6e3dd2e4a..fa1c6ccd6 100644
--- a/dask/array/_shuffle.py
+++ b/dask/array/_shuffle.py
@@ -1,7 +1,9 @@
 from __future__ import annotations
 
 import copy
+import math
 from itertools import count, product
+from typing import Literal
 
 import numpy as np
 
@@ -13,7 +15,7 @@ from dask.base import tokenize
 from dask.highlevelgraph import HighLevelGraph
 
 
-def shuffle(x, indexer: list[list[int]], axis):
+def shuffle(x, indexer: list[list[int]], axis: int, chunks: Literal["auto"] = "auto"):
     """
     Reorders one dimensions of a Dask Array based on an indexer.
 
@@ -34,6 +36,12 @@ def shuffle(x, indexer: list[list[int]], axis):
         each group will end up in exactly one chunk.
     axis: int
         The axis to shuffle along.
+    chunks: "auto"
+        Hint on how to rechunk if single groups are becoming too large. The default is
+        to split chunks along the other dimensions evenly to keep the chunksize
+        consistent. The rechunking is done in a way that ensures that non all-to-all
+        network communication is necessary, chunks are only split and not combined with
+        other chunks.
 
     Examples
     --------
@@ -66,6 +74,9 @@ def shuffle(x, indexer: list[list[int]], axis):
             f"Shuffling only allowed with known chunk sizes. {unknown_chunk_message}"
         )
     assert isinstance(axis, int), "axis must be an integer"
+    _validate_indexer(x.chunks, indexer, axis)
+
+    x = _rechunk_other_dimensions(x, max(map(len, indexer)), axis, chunks)
 
     token = tokenize(x, indexer, axis)
     out_name = f"shuffle-{token}"
@@ -79,7 +90,54 @@ def shuffle(x, indexer: list[list[int]], axis):
     return Array(graph, out_name, chunks, meta=x)
 
 
-def _shuffle(chunks, indexer, axis, in_name, out_name, token):
+def _rechunk_other_dimensions(
+    x: Array, longest_group: int, axis: int, chunks: Literal["auto"]
+) -> Array:
+    assert chunks == "auto", "Only auto is supported for now"
+    chunksize_tolerance = config.get("array.shuffle.chunksize-tolerance")
+
+    if longest_group <= max(x.chunks[axis]) * chunksize_tolerance:
+        # We are staying below our threshold, so don't rechunk
+        return x
+
+    # How much the chunksizes on our shuffle axis will increase
+    chunksize_inc_factor = longest_group / max(x.chunks[axis])
+
+    new_chunks = []
+    for i in range(len(x.chunks)):
+        if i == axis:
+            new_chunks.append(x.chunks[i])
+            continue
+
+        new_chunksizes = []
+        # calculate what the max chunk size in this dimension is and split every
+        # chunk that is larger than that. We split the increase factor evenly
+        # between all dimensions that are not shuffled.
+        up_chunksize_limit_for_dim = max(x.chunks[i]) / (
+            chunksize_inc_factor ** (1 / (len(x.chunks) - 1))
+        )
+        for c in x.chunks[i]:
+            if c > chunksize_tolerance * up_chunksize_limit_for_dim:
+                factor = math.ceil(c / up_chunksize_limit_for_dim)
+
+                # Ensure that we end up at least with chunksize 1
+                factor = min(factor, c)
+
+                chunksize, remainder = divmod(c, factor)
+                nc = [chunksize] * factor
+                for ii in range(remainder):
+                    # Add remainder parts to the first few chunks
+                    nc[ii] += 1
+                new_chunksizes.extend(nc)
+
+            else:
+                new_chunksizes.append(c)
+        new_chunks.append(tuple(new_chunksizes))
+
+    return x.rechunk(tuple(new_chunks))
+
+
+def _validate_indexer(chunks, indexer, axis):
     if not isinstance(indexer, list) or not all(isinstance(i, list) for i in indexer):
         raise ValueError("indexer must be a list of lists of positional indices")
 
@@ -93,6 +151,10 @@ def _shuffle(chunks, indexer, axis, in_name, out_name, token):
             f"Indexer contains out of bounds index. Dimension only has {sum(chunks[axis])} elements."
         )
 
+
+def _shuffle(chunks, indexer, axis, in_name, out_name, token):
+    _validate_indexer(chunks, indexer, axis)
+
     if len(indexer) == len(chunks[axis]):
         # check if the array is already shuffled the way we want
         ctr = 0
diff --git a/dask/array/core.py b/dask/array/core.py
index 3009dcb11..a35c05f76 100644
--- a/dask/array/core.py
+++ b/dask/array/core.py
@@ -17,7 +17,7 @@ from itertools import product, zip_longest
 from numbers import Integral, Number
 from operator import add, mul
 from threading import Lock
-from typing import Any, TypeVar, Union, cast
+from typing import Any, Literal, TypeVar, Union, cast
 
 import numpy as np
 from numpy.typing import ArrayLike
@@ -2765,7 +2765,8 @@ class Array(DaskMethodsMixin):
     def shuffle(
         self,
         indexer: list[list[int]],
-        axis,
+        axis: int,
+        chunks: Literal["auto"] = "auto",
     ):
         """Reorders one dimensions of a Dask Array based on an indexer.
 
@@ -2777,7 +2778,7 @@ class Array(DaskMethodsMixin):
         """
         from dask.array._shuffle import shuffle
 
-        return shuffle(self, indexer, axis)
+        return shuffle(self, indexer, axis, chunks)
 
     @property
     def real(self):
diff --git a/dask/array/slicing.py b/dask/array/slicing.py
index a1bf99e85..c18e52939 100644
--- a/dask/array/slicing.py
+++ b/dask/array/slicing.py
@@ -565,8 +565,9 @@ def take(outname, inname, chunks, index, axis=0):
 
     if not np.isnan(chunks[axis]).any():
         from dask.array._shuffle import _shuffle
+        from dask.array.utils import arange_safe, asarray_safe
 
-        arange = np.arange(np.sum(chunks[axis]))
+        arange = arange_safe(np.sum(chunks[axis]), like=index)
         if len(index) == len(arange) and np.abs(index - arange).sum() == 0:
             # TODO: This should be a real no-op, but the call stack is
             # too deep to do this efficiently for now
@@ -577,7 +578,7 @@ def take(outname, inname, chunks, index, axis=0):
         average_chunk_size = int(sum(chunks[axis]) / len(chunks[axis]))
 
         indexer = []
-        index = np.asarray(index)
+        index = asarray_safe(index, like=index)
         for i in range(0, len(index), average_chunk_size):
             indexer.append(index[i : i + average_chunk_size].tolist())
 


Test patch:
diff --git a/dask/array/tests/test_shuffle.py b/dask/array/tests/test_shuffle.py
index 2882cbc87..86cadcfc2 100644
--- a/dask/array/tests/test_shuffle.py
+++ b/dask/array/tests/test_shuffle.py
@@ -6,6 +6,7 @@ import pytest
 import dask
 import dask.array as da
 from dask.array import assert_eq, shuffle
+from dask.array._shuffle import _rechunk_other_dimensions
 from dask.core import flatten
 
 
@@ -20,19 +21,19 @@ def darr(arr):
 
 
 @pytest.mark.parametrize(
-    "indexer, chunks",
+    "indexer, chunks, other_chunks",
     [
-        ([[1, 5, 6], [0, 2, 3, 4, 7]], (3, 5)),
-        ([[1, 5, 6], [0, 3], [4, 2, 7]], (5, 3)),
-        ([[1], [0, 6, 5, 3, 2, 4], [7]], (1, 6, 1)),
-        ([[1, 5, 1, 5, 1, 5], [1, 6, 4, 2, 7]], (6, 5)),
+        ([[1, 5, 6], [0, 2, 3, 4, 7]], (3, 5), (2, 1)),
+        ([[1, 5, 6], [0, 3], [4, 2, 7]], (5, 3), (2, 1)),
+        ([[1], [0, 6, 5, 3, 2, 4], [7]], (1, 6, 1), (1, 1, 1)),
+        ([[1, 5, 1, 5, 1, 5], [1, 6, 4, 2, 7]], (6, 5), (1, 1, 1)),
     ],
 )
-def test_shuffle(arr, darr, indexer, chunks):
+def test_shuffle(arr, darr, indexer, chunks, other_chunks):
     result = darr.shuffle(indexer, axis=1)
     expected = arr[:, list(flatten(indexer))]
     assert_eq(result, expected)
-    assert result.chunks[0] == darr.chunks[0]
+    assert result.chunks[0] == other_chunks
     assert result.chunks[1] == chunks
 
 
@@ -94,3 +95,24 @@ def test_shuffle_no_op_with_correct_indexer():
     result = arr.shuffle(indexer, axis=0)
     assert result.dask == arr.dask
     assert_eq(arr, result)
+
+
+def test_resize_other_dimensions():
+    arr = da.random.random((250, 50), chunks=((45, 100, 38, 67), 10))
+    result = _rechunk_other_dimensions(arr, 20, 1, "auto")
+    assert result.chunks == ((45, 50, 50, 38, 34, 33), (10,) * 5)
+    assert_eq(arr, result)
+
+    arr = da.random.random((250, 50, 20), chunks=((45, 100, 38, 67), 10, 10))
+    result = _rechunk_other_dimensions(arr, 20, 1, "auto")
+    assert result.chunks == ((45, 50, 50, 38, 67), (10,) * 5, (5,) * 4)
+    assert_eq(arr, result)
+
+    result = _rechunk_other_dimensions(arr, 40, 1, "auto")
+    assert result.chunks == ((45, 50, 50, 38, 34, 33), (10,) * 5, (5,) * 4)
+    assert_eq(arr, result)
+
+    arr = da.random.random((250, 50, 5), chunks=((45, 100, 38, 67), 10, 1))
+    result = _rechunk_other_dimensions(arr, 40, 1, "auto")
+    assert result.chunks == ((45, 50, 50, 38, 34, 33), (10,) * 5, (1,) * 5)
+    assert_eq(arr, result)
diff --git a/dask/dataframe/tests/test_dataframe.py b/dask/dataframe/tests/test_dataframe.py
index f0f65f72d..a2cc707c2 100644
--- a/dask/dataframe/tests/test_dataframe.py
+++ b/dask/dataframe/tests/test_dataframe.py
@@ -3610,6 +3610,7 @@ def test_cov_series():
 
 
 @pytest.mark.gpu
+@pytest.mark.skip(reason="https://github.com/rapidsai/cudf/issues/16560")
 @pytest.mark.parametrize(
     "numeric_only",
     [None, True, False],
@@ -3682,6 +3683,7 @@ def test_corr():
 
 
 @pytest.mark.gpu
+@pytest.mark.skip(reason="https://github.com/rapidsai/cudf/issues/16560")
 def test_corr_gpu():
     cudf = pytest.importorskip("cudf")