Repository: dask/dask

Issue:
High memory usage for reductions.median but not reductions.mean
I'm attempting to median stack 200 x 4k x 4k images. I've been having trouble with the memory allocation so I tried to simplify my problem by going to 2d and was surprised to see that doing a median reduction can take 50 GB of memory and crash python.

This code works as expected in ipython. Almost no memory is used and this completes in less than a minute.
```python
import dask.array as da
x = da.random.normal(10, 0.1, size=(100000, 100000), chunks=(10, 100000))
y = da.reductions.mean(x, axis=1)
%time z = y.compute()
```
`CPU times: user 6min 38s, sys: 2.36 s, total: 6min 41s
Wall time: 56.3 s`

Running the following code causes python to crash after activity monitor showing it using more than 50 GB of RAM.

```python
import dask.array as da
x = da.random.normal(10, 0.1, size=(100000, 100000), chunks=(10, 100000))
y = da.reductions.median(x, axis=1)
%time z = y.compute()
```
![Screen Shot 2021-06-09 at 4 32 29 PM](https://user-images.githubusercontent.com/5331996/121426745-c66f7480-c928-11eb-8f5e-17bdf4c69a48.png)

![Screen Shot 2021-06-09 at 4 47 17 PM](https://user-images.githubusercontent.com/5331996/121427191-45fd4380-c929-11eb-8bdc-5ed0985ec02d.png)

**Anything else we need to know?**:
I did also try this in regular python (as opposed to ipython) and got the same crash. The odd chunk size is an attempt to make sure each chunk has the full array necessary to compute the median in one chunk.



**Environment**:

- Dask version: 2021.05.1
- Python version: 3.8.10
- IPython version: 7.22.0
- Operating System: MacOS 11.3.1
- Install method (conda, pip, source): conda



Discussion:
Mean works because the data creation and initial data reduction steps are
fused, so we create and destroy data at the same pace.

In the median case there is an intervening step, where we rechunk the array
along the axis to be reduced over to ensure that we can do the median (we
don't have a parallel median algorithm).  This creates a bit of a
disconnect in the graph, and so the worker creates some data and then
doesn't immediately consume that data with the first pass of the
reduction.  Instead it asks the scheduler what's next (in the distributed
case) or waits for neighboring data creation chunks to be done as well for
the eventual rechunk.  The workers are eating themselves to death :)

The solution to this problem is probably to rechunk less aggressively. In
your case we don't even need to rechunk at all.  We can fix this by
changing line 1586 here:

https://github.com/dask/dask/blob/eed298b2335401d2cb8e85ebb44888eb59937cf8/dask/array/reductions.py#L1569-L1588

Probably we should

   1. Only call rechunk if any(x.numblocks[ax] > 1 for ax in axis)
   2. Be careful about using "auto" on all other dimensions, and instead
   make sure that we only reduce the size of chunks, rather than increase them
   / glom them together (which can cause scheduling artifacts)

The first task is the easiest and would solve your problem immediately.
It's a decent task for anyone with mild experience with dask array.  The
second task would be good too, but is not necessary.

On Wed, Jun 9, 2021 at 3:48 PM Curtis McCully ***@***.***>
wrote:

> I'm attempting to median stack 200 x 4k x 4k images. I've been having
> trouble with the memory allocation so I tried to simplify my problem by
> going to 2d and was surprised to see that doing a median reduction can take
> 50 GB of memory and crash python.
>
> This code works as expected in ipython. Almost no memory is used and this
> completes in less than a minute.
>
> import dask.array as dax = da.random.normal(10, 0.1, size=(100000, 100000), chunks=(10, 100000))y = da.reductions.mean(x, axis=1)%time z = y.compute()
>
> CPU times: user 6min 38s, sys: 2.36 s, total: 6min 41s Wall time: 56.3 s
>
> Running the following code causes python to crash after activity monitor
> showing it using more than 50 GB of RAM.
>
> import dask.array as dax = da.random.normal(10, 0.1, size=(100000, 100000), chunks=(10, 100000))y = da.reductions.median(x, axis=1)%time z = y.compute()
>
> [image: Screen Shot 2021-06-09 at 4 32 29 PM]
> <https://user-images.githubusercontent.com/5331996/121426745-c66f7480-c928-11eb-8f5e-17bdf4c69a48.png>
>
> [image: Screen Shot 2021-06-09 at 4 47 17 PM]
> <https://user-images.githubusercontent.com/5331996/121427191-45fd4380-c929-11eb-8bdc-5ed0985ec02d.png>
>
> *Anything else we need to know?*:
> I did also try this in regular python (as opposed to ipython) and got the
> same crash. The odd chunk size is an attempt to make sure each chunk has
> the full array necessary to compute the median in one chunk.
>
> *Environment*:
>
>    - Dask version: 2021.05.1
>    - Python version: 3.8.10
>    - IPython version: 7.22.0
>    - Operating System: MacOS 11.3.1
>    - Install method (conda, pip, source): conda
>
> —
> You are receiving this because you are subscribed to this thread.
> Reply to this email directly, view it on GitHub
> <https://github.com/dask/dask/issues/7780>, or unsubscribe
> <https://github.com/notifications/unsubscribe-auth/AACKZTAERPT66R2TZZLG6Q3TR7HRPANCNFSM46MZLC5A>
> .
>


Resolving patch:
diff --git a/dask/array/reductions.py b/dask/array/reductions.py
--- a/dask/array/reductions.py
+++ b/dask/array/reductions.py
@@ -1569,7 +1569,7 @@ def trace(a, offset=0, axis1=0, axis2=1, dtype=None):
 @derived_from(np)
 def median(a, axis=None, keepdims=False, out=None):
     """
-    This works by automatically chunking the reduced axes to a single chunk
+    This works by automatically chunking the reduced axes to a single chunk if necessary
     and then calling ``numpy.median`` function across the remaining dimensions
     """
     if axis is None:
@@ -1583,7 +1583,9 @@ def median(a, axis=None, keepdims=False, out=None):
 
     axis = [ax + a.ndim if ax < 0 else ax for ax in axis]
 
-    a = a.rechunk({ax: -1 if ax in axis else "auto" for ax in range(a.ndim)})
+    # rechunk if reduced axes are not contained in a single chunk
+    if builtins.any(a.numblocks[ax] > 1 for ax in axis):
+        a = a.rechunk({ax: -1 if ax in axis else "auto" for ax in range(a.ndim)})
 
     result = a.map_blocks(
         np.median,
@@ -1616,7 +1618,9 @@ def nanmedian(a, axis=None, keepdims=False, out=None):
 
     axis = [ax + a.ndim if ax < 0 else ax for ax in axis]
 
-    a = a.rechunk({ax: -1 if ax in axis else "auto" for ax in range(a.ndim)})
+    # rechunk if reduced axes are not contained in a single chunk
+    if builtins.any(a.numblocks[ax] > 1 for ax in axis):
+        a = a.rechunk({ax: -1 if ax in axis else "auto" for ax in range(a.ndim)})
 
     result = a.map_blocks(
         np.nanmedian,


Test patch:
diff --git a/dask/array/tests/test_reductions.py b/dask/array/tests/test_reductions.py
--- a/dask/array/tests/test_reductions.py
+++ b/dask/array/tests/test_reductions.py
@@ -702,6 +702,22 @@ def test_median(axis, keepdims, func):
     )
 
 
+@pytest.mark.parametrize("func", ["median", "nanmedian"])
+@pytest.mark.parametrize("axis", [0, [0, 2], 1])
+def test_median_does_not_rechunk_if_whole_axis_in_one_chunk(axis, func):
+    x = np.arange(100).reshape((2, 5, 10))
+    d = da.from_array(x, chunks=(2, 1, 10))
+
+    actual = getattr(da, func)(d, axis=axis)
+    expected = getattr(np, func)(x, axis=axis)
+    assert_eq(actual, expected)
+    does_rechunk = "rechunk" in str(dict(actual.__dask_graph__()))
+    if axis == 1:
+        assert does_rechunk
+    else:
+        assert not does_rechunk
+
+
 @pytest.mark.parametrize("method", ["sum", "mean", "prod"])
 def test_object_reduction(method):
     arr = da.ones(1).astype(object)