Repository: dask/dask

Issue:
Add scheduler option to `assert_eq` utilities
In https://github.com/dask/dask/pull/8559 we added a `scheduler=` keyword argument to `dask.dataframe.assert_eq` to allow for cases where we want to run the computes that happen in `assert_eq` on a specified scheduler. As a follow-up, we should add a similar `scheduler=` keyword argument to our other `assert_eq` utilities in `dask.array`

https://github.com/dask/dask/blob/a504115312b1e591604276a3098b2e0076bf49e1/dask/array/utils.py#L261

and `dask.bag`

https://github.com/dask/dask/blob/a504115312b1e591604276a3098b2e0076bf49e1/dask/bag/utils.py#L1


Discussion:

Patch:
diff --git a/dask/array/utils.py b/dask/array/utils.py
--- a/dask/array/utils.py
+++ b/dask/array/utils.py
@@ -226,7 +226,9 @@ def _check_chunks(x):
     return x
 
 
-def _get_dt_meta_computed(x, check_shape=True, check_graph=True, check_chunks=True):
+def _get_dt_meta_computed(
+    x, check_shape=True, check_graph=True, check_chunks=True, scheduler=None
+):
     x_original = x
     x_meta = None
     x_computed = None
@@ -240,7 +242,7 @@ def _get_dt_meta_computed(x, check_shape=True, check_graph=True, check_chunks=Tr
         if check_chunks:
             # Replace x with persisted version to avoid computing it twice.
             x = _check_chunks(x)
-        x = x.compute(scheduler="sync")
+        x = x.compute(scheduler=scheduler)
         x_computed = x
         if hasattr(x, "todense"):
             x = x.todense()
@@ -266,6 +268,7 @@ def assert_eq(
     check_meta=True,
     check_chunks=True,
     check_type=True,
+    scheduler="sync",
     **kwargs,
 ):
     a_original = a
@@ -277,10 +280,18 @@ def assert_eq(
         b = np.array(b)
 
     a, adt, a_meta, a_computed = _get_dt_meta_computed(
-        a, check_shape=check_shape, check_graph=check_graph, check_chunks=check_chunks
+        a,
+        check_shape=check_shape,
+        check_graph=check_graph,
+        check_chunks=check_chunks,
+        scheduler=scheduler,
     )
     b, bdt, b_meta, b_computed = _get_dt_meta_computed(
-        b, check_shape=check_shape, check_graph=check_graph, check_chunks=check_chunks
+        b,
+        check_shape=check_shape,
+        check_graph=check_graph,
+        check_chunks=check_chunks,
+        scheduler=scheduler,
     )
 
     if str(adt) != str(bdt):
diff --git a/dask/bag/utils.py b/dask/bag/utils.py
--- a/dask/bag/utils.py
+++ b/dask/bag/utils.py
@@ -1,7 +1,7 @@
-def assert_eq(a, b):
+def assert_eq(a, b, scheduler="sync"):
     if hasattr(a, "compute"):
-        a = a.compute(scheduler="sync")
+        a = a.compute(scheduler=scheduler)
     if hasattr(b, "compute"):
-        b = b.compute(scheduler="sync")
+        b = b.compute(scheduler=scheduler)
 
     assert a == b


Test patch:
diff --git a/dask/array/tests/test_array_utils.py b/dask/array/tests/test_array_utils.py
--- a/dask/array/tests/test_array_utils.py
+++ b/dask/array/tests/test_array_utils.py
@@ -2,7 +2,9 @@
 import pytest
 
 import dask.array as da
+from dask.array.core import Array
 from dask.array.utils import assert_eq, meta_from_array
+from dask.local import get_sync
 
 asarrays = [np.asarray]
 
@@ -87,3 +89,28 @@ def test_meta_from_array_type_inputs():
 def test_assert_eq_checks_dtype(a, b):
     with pytest.raises(AssertionError, match="a and b have different dtypes"):
         assert_eq(a, b)
+
+
+@pytest.mark.parametrize(
+    "a,b",
+    [
+        (1.0, 1.0),
+        ([1, 2], [1, 2]),
+        (da.array([1, 2]), da.array([1, 2])),
+    ],
+)
+def test_assert_eq_scheduler(a, b):
+    counter = 0  # Counts how many times `custom_scheduler` is executed.
+
+    def custom_scheduler(*args, **kwargs):
+        nonlocal counter
+        counter += 1
+        return get_sync(*args, **kwargs)
+
+    assert_eq(a, b)
+    assert counter == 0
+
+    assert_eq(a, b, scheduler=custom_scheduler)
+    # `custom_scheduler` should be executed twice if `a` and `b` are Arrays.
+    n_da_arrays = len([x for x in [a, b] if isinstance(x, Array)])
+    assert counter == n_da_arrays