map_overlap performance on large inputs

[emilmelnikov]

We had a discussion with @GenevieveBuckley and @jni about our attempts to use dask arrays in ilastik.

Our application is a 5D image viewer that also interactively applies ML algorithms to user data. One of our workflows needs to compute various filters from the input data of arbitrary size. For the purposes of this discussion, assume that the input data is blocked (HDF5, Zarr, N5). All these filters are convolutional (Gaussian, DoG etc.), so they can be computed only for blocks that fully cover the current GUI viewport. Also, in order to maintain an acceptable latency, and to reduce the amount of wasted compute, block sizes should be quite small (we use 64^3 for 3D data). This creates a performance problem if one wants to use dask arrays for all lazy computations in a pipeline, due to the large size of a dask graph for intermediate computations.

import contextlib
import time

import dask
import dask.array
import zarr


def identity(x):
    return x


@contextlib.contextmanager
def measure(label):
    t0 = time.perf_counter()
    yield
    t = time.perf_counter() - t0
    print(f"{label}: {t:.9} seconds")


print(f"{dask.__version__ = }")

shape = (2**11,) * 3
chunks = (64,) * 3
src = zarr.zeros(shape, chunks=chunks, dtype="uint8")
print(f"{src.shape = }")
print(f"{src.chunks = }")
print(f"{src.nbytes / (2**30) = }")

with measure("create dask array 'a' from zarr array 'src'"):
    a = dask.array.from_zarr(src)
print(f"{len(a.__dask_graph__()) = }")

with measure("map_blocks dask array 'a' to dask array 'b'"):
    b = dask.array.map_blocks(identity, a)
print(f"{len(b.__dask_graph__()) = }")

with measure("map_overlap dask array 'a' to dask array 'c'"):
    c = dask.array.map_overlap(identity, a, depth=1, boundary="none")
print(f"{len(c.__dask_graph__()) = }")

Output:

dask.__version__ = '2021.03.0'
src.shape = (2048, 2048, 2048)
src.chunks = (64, 64, 64)
src.nbytes / (2**30) = 8.0
create dask array 'a' from zarr array 'src': 0.021725723 seconds
len(a.__dask_graph__()) = 32769
map_blocks dask array 'a' to dask array 'b': 0.000862336 seconds
len(b.__dask_graph__()) = 65537
map_overlap dask array 'a' to dask array 'c': 8.43978025 seconds
len(c.__dask_graph__()) = 994425

I wonder if #6438 is somehow related to this problem?

[mrocklin]

cc @rjzamora and @ian-r-rose who have thought about high level graphs for arrays. One way to solve this problem would be to make a high level graph layer for map-overlap and then implement culling/dependency for that layer.

[rjzamora]

Once #7381 is in, it shouldn't be too difficult to pull most array creation/io under the Blockwise umbrella. From an initial glance at the example above, it seems that a reasonable first step is to convert from_array and from_zarr into Blockwise layers (to avoid materializing a large graph during the creation step).

[GenevieveBuckley]

Looks like #7381 was merged 6 hours ago, so great timing! @rjzamora would you be willing to review a PR converting from_array and from_zarr into Blockwise?

[rjzamora]

@rjzamora would you be willing to review a PR converting from_array and from_zarr into Blockwise?

Absolutely happy to help in any way I can! Note that I have started to migrate some DataFrame-specific IO in #7415, but I haven't tackled DataFrame.from_array yet (which I would expect to require similar logic). Also, cc @ian-r-rose, who has been thinking about Blockwise IO in dask.array lately.

[GenevieveBuckley]

Started looking at how we could make a blockwise from_array this morning, and have only just realised #7417 was opened over the weekend (i.e. Friday for the USA). Going to take a look at that now.

[GenevieveBuckley]

Spent some more time looking at map_overlap today, it seems that task graphs get very complicated very quickly. I thought a toy example might help me:

import dask.array as da

def myfunc(x):
    return x + 1

arr = da.from_array(np.ones((2, 4)).astype(int), chunks=2)
arr

result_from_blocks = arr.map_blocks(myfunc)
result_from_blocks.visualize()

result_with_overlap = arr.map_overlap(myfunc, depth=1, align_arrays=False, trim=False, boundary=None)
result_with_overlap.visualize()

This task graph is a bit too complicated to be very useful. I went down a bit of a rabbit hole looking for tools that would help me compare graphviz graphs more easily, but I suspect it's better to try a different approach for working out what's going on.

[mrocklin]

The blockwise work will be great in accelerating things like from_array or from_zarr, however it won't help with map_overlap. Instead, we'll need to repeat the treatment that we did for blockwise on map_overlap. We'll probably need to make a high level layer for map_overlap (or the various sub-components of that operation) to make this work well. I would expect this to be at least a week of work. from someone.

[GenevieveBuckley]

This is a naive question perhaps, but what is the reason overlap_internal can't be a slice operation?

[GenevieveBuckley]

FYI @gjoseph92 & @ian-r-rose

My colleague just pointed out a spot where re-inventing a version of map_overlap helped avoid some of the performance problems we've been talking about:
https://github.com/royerlab/dexp/blob/master/dexp/processing/utils/scatter_gather_i2i.py

[GenevieveBuckley]

@jni requested links to:

• the map_overlap PR: A HighLevelGraph abstract layer for map_overlap #7595
• the slicing PR: Array slicing HighLevelGraph layer #7655
• this blogpost: https://blog.dask.org/2021/07/07/high-level-graphs

[GenevieveBuckley]

Ian's written a summary of the current status and next steps necessary for the high level graph work for slicing and array overlap. The gist is here: https://gist.github.com/ian-r-rose/4221ebf52f3423203640c498fb815f21