Repository: dask/dask

Issue:
Dask array rechunk should do intermediate sized rechunking when appropriate
Consider rechunking from chunks=(1, 100) to chunks=(100, 1) on an array of shape (100, 100).

A direct rechunk operation creates ~10,000 tasks corresponding to intermediate arrays of shape (1,1).

However, if we rechunk first to shape (10,10) (the same array size as all of our chunks) and then to our final chunk shape, there would be only ~2000 tasks, because the smallest intermediate arrays have 10 elements.

I don't know what the general rules are for how to optimize rechunkings but this makes a big difference in computational efficiency -- O(n^2) vs O(n^1.5) for square arrays with a side of length n. For n=10,000, this is a ~50x speed up (assuming that the task overhead dominates for small tasks).

Credit to @freeman-lab for this idea, which apparently works really nicely in Thunder. Jeremy, if you could share how you calculate intermediate chunks sizes that would be helpful!



Discussion:
I've run into this with other ndarray-on-spark systems and have heard that it performs decently.

@shoyer from your description I would propose:
 given final chunk size  `n = final_chunk.size`, and the array has `ndim` number of dimensions. Then the intermediate chunks can be square, where each dimension's length is the _ndim_th root of _n_. Like `intermediate_chunk = (floor(n**(1./ndim)),) * ndim`

But I'm not familiar enough with the rechunking internals to know if this is too naive.

If you do the rechunking repeatedly, each time scaling the longer chunk axis by a factor of two, you can get the number of intermediate chunks created down to O(n log(n)). OTOH, I have no idea how this scales in practice.

@shoyer I'm not sure I understand. Could you give an example?

Again, suppose you have an n x n array completely chunked along one axis with chunks of size (n, 1). In the first chunking operation, create chunks of size (n/2, 2). This involves 2n intermediate arrays. Then repeat the process log(n) times, and you have your final desired chunk size of (1, n). This involves only O(n log(n)) tasks. But you also created O(log(n)) temporary copies of the full dataset.

On Sat, Jul 18, 2015 at 5:34 PM, Blake Griffith notifications@github.com
wrote:

> ## @shoyer I'm not sure I understand. Could you give an example?
> 
> Reply to this email directly or view it on GitHub:
> https://github.com/ContinuumIO/dask/issues/417#issuecomment-122608426

I did a few benchmarks my machine:
http://nbviewer.ipython.org/gist/shoyer/425cb3dc9101c235fe86

For 1024 x 1024 arrays, I get a 15x speedup using one intermediate rechunking, and a 30x speedup using 3 intermediates -- bringing the calculation down from 4 minutes to 7 seconds. So it looks like a pretty clear win.

This is very cool @shoyer ! The size estimation we've done for this kind of chunking has been more about optimizing for axis swaps in large shuffles, and in that case chunks of approximately 100-150mb have empirically proven optimal, in that smaller chunks are slower (due to increased shuffle complexity) and larger chunks run into memory issues.

The optimal sizes here clearly have different constraints, but cool that the concept proves useful! 

@freeman-lab Do you guys use multiple intermediates in axis swaps or just one?

@shoyer just one... at least in Spark I'm fairly certain the time incurred just from processing the additional shuffles would outweigh any other gains, but I don't have hard numbers.

Btw, we've just about finished moving all the custom chunking logic into Bolt (https://github.com/bolt-project/bolt/pull/67), if you want to check out more of the details there.

@shoyer how do these speeds compare to just throwing the array into an empty 1024x1024 numpy array and then pulling an array with different chunks back out?  How would your benchmarks change with an array that didn't fit in memory?  Are the current scheduling routines smart enough to manage a reshuffle like this without filling up the cache and writing back to disk?

> how do these speeds compare to just throwing the array into an empty 1024x1024 numpy array and then pulling an array with different chunks back out?

I'm not quite sure what you're getting at here -- I think we can safely presume that creating the 1024x1024 numpy array directly would be far, far faster.

> How would your benchmarks change with an array that didn't fit in memory?

It takes 53 seconds merely to compute the _graph_ to do the naive rechunking, before any of the values are computed. Naive rechunking will be entirely unworkable for arrays that don't fit in memory.

> Are the current scheduling routines smart enough to manage a reshuffle like this without filling up the cache and writing back to disk?

My sense is that doing a full rechunking will pretty much require spilling to disk, regardless -- or distributing an array in memory across many machines.

> I'm not quite sure what you're getting at here -- I think we can safely presume that creating the 1024x1024 numpy array directly would be far, far faster.
> 
> My sense is that doing a full rechunking will pretty much require spilling to disk, regardless -- or distributing an array in memory across many machines.

In the single machine case if this requires spilling to disk and if creating the temporary explicitly is faster then maybe the optimal approach is to explicitly store the result into a shared array and then pull out from it again.

This issue has become a bit more interesting now with distributed scheduling.  I plan to run a benchmark on a few approaches soon.

Verified that in practice with the distributed scheduler the log-based rechunking seems nicer.  I also found that an explicit `rebalance` operation after rechunking was cheap and useful.  Things tend to get a bit unbalanced during computation.  I have yet to try this on a large cluster.  

It's possible that we'll be scheduler bound.

@freeman-lab roughly what kind of performance does Thunder get on Spark for tranposing the direction of partitioning?  Do you have ballpark numbers for a common dataset shape that you often use?

@mrocklin we're putting some numbers together now with recent versions of everything, will post here in just a bit

@freeman-lab even order-of-magnitude numbers would be really valuable here.  Something along the lines of the following would be really helpful.
-  Our data is usually shaped like 10000 by 10000000
-  A transpose of chunk shape from (10000, 1) to (1, 10000000) usually takes around 10s-2m on 20 m3.2xlarges

At the moment I'm currently curious if I'm anywhere in the same ballpark or not.  I'm trying to figure out how far I can get without building an explicit shuffle operation and just using task scheduling throughout.  My hope is to get a quick "nope, definitely not close" so that I know if I should continue to pursue my current path.

cc @broxtronix in case he has ballpark numbers 

@mrocklin cool gotcha so a typical medium-sized dataset for us might be (10000, 512, 512, 4), that's around ~20 GB as `uint16` and we'll often go from 10000 chunks of size (1, 1000000) to 1000000 chunks of size (10000, 1). The complete operation of loading from disk and doing the rechunking on 10 nodes might take ~2 minutes. Those numbers are from an internal cluster, which based on other benchmarking we've done might be 2-4x faster than the generic `m3`, but the exact factor there depends a lot on whether it's compute / network / io bound.

Is there a sequence of commands I can run in Thunder to do this on random data?

Do you have multiple chunks per spark partition?

yup yup, with v1.0.0 it's as easy as

```
import thunder as td

data = td.images.fromrandom(shape=(10000,512,512,4), npartitions=320, engine=sc)
data.toseries().compute()
```

where `sc` is a valid `SparkContext`

We typically do have multiple chunks per spark partition, assuming we set the partitioning to be roughly 2-3x the number of cores, which generally gives the best performance.

We just confirmed the 2 minutes is correct for that size on 10 nodes, also confirmed that doing the same thing for `(2000, 1000, 2000, 30)` data, going from `(1,60000000)` chunks to `(2000,1)` chunks is about 10 minutes on 20 nodes.

cc @jwittenbach who's been running these tests

Thank you.  These numbers are great.  Will crunch away on dask stuff and see how close I can get.

cool, very curious to hear how it goes!

OK, using a cascading transpose as suggested by @shoyer we can flip around the 20GB array in roughly two minutes.  

https://gist.github.com/mrocklin/01f39076b86fe73c71c353be830e4d0f

Looking at the Dask web UI this _does_ seem to take a fair amount of RAM.  It doesn't appear to be particularly sensitive to how many steps we use in the transpose-cascade.  Using only every second or third step or using significantly more partitions results in only a marginal increase (around 50%).  The scheduler is nowhere near peak when doing these operations, so we should be able to scale quite a bit larger.

The interesting thing we've learned here is that, as long as we keep the scheduler fast and pay attention to batching, it seems to be reasonable to avoid building a separate shuffle operation, at least for now.

@shoyer this might interest you: https://www.youtube.com/watch?v=W_VBpv27soU&feature=youtu.be

@mrocklin Yes, interesting to see this working! I would definitely try to benchmark against Bolt on the same cluster. I would not be surprised at all if EC2 has significantly worse networking than @freeman-lab's cluster, even with enhanced networking.

I can get this down to three minutes by splitting with a finer resolution, `2**5` rather than `2**2`.  This shifts burden from the network to the scheduler (fewer copies, more tasks) but the scheduler seems to be able to handle it just fine.  

I tested network bandwidth by trying `(x + x.T)` on a large square matrix and got up to around 300 MB/s steady output averaged across the network.

@mrocklin Something is wrong with the formatting on your last comment. 

Fixed

By resolution you mean the step size for each rechunk? That makes sense. 

Yeah, this chunks more finely and so accomplishes the movement in fewer stages with more tasks per stage.  We're nowhere near saturating the scheduler though (which is now quite speedy) so all is well.

Changed a couple things in the scheduler gets us down to around 2 minutes.  

Also, I realized that my test cluster didn't have any compression libraries installed.  After conda installing lz4 doing the transpose in more steps wasn't as bad, and came down to about 2.5 minutes.

In the two-step transpose I'm nowhere near saturating the bandwidth of my network.  The scheduler also doesn't appear to be particularly burdened.  Not sure where the current bottleneck is.

Actually, I just measured bandwidth by doing `x + x.T` on a large square matrix and came up with the same bandwidth measurement.  It could be that EC2 is a bit bogged down at the moment.  It would be interesting to run the above computation on a dedicated cluster.


Resolving patch:
diff --git a/dask/array/rechunk.py b/dask/array/rechunk.py
--- a/dask/array/rechunk.py
+++ b/dask/array/rechunk.py
@@ -9,7 +9,7 @@
 
 import heapq
 
-from itertools import product, chain
+from itertools import product, chain, count
 from operator import getitem, add, mul, itemgetter
 
 import numpy as np
@@ -53,12 +53,12 @@ def _intersect_1d(breaks):
     >>> old = cumdims_label(((2, 2, 1), (5,)), 'o')
 
     >>> _intersect_1d(_breakpoints(old[0], new[0]))  # doctest: +NORMALIZE_WHITESPACE
-    (((0, slice(0, 2, None)),),
-     ((1, slice(0, 2, None)), (2, slice(0, 1, None))))
+    [[(0, slice(0, 2, None))],
+     [(1, slice(0, 2, None)), (2, slice(0, 1, None))]]
     >>> _intersect_1d(_breakpoints(old[1], new[1]))  # doctest: +NORMALIZE_WHITESPACE
-    (((0, slice(0, 2, None)),),
-     ((0, slice(2, 4, None)),),
-     ((0, slice(4, 5, None)),))
+    [[(0, slice(0, 2, None))],
+     [(0, slice(2, 4, None))],
+     [(0, slice(4, 5, None))]]
 
     Parameters
     ----------
@@ -75,11 +75,14 @@ def _intersect_1d(breaks):
     last_end = 0
     old_idx = 0
     ret = []
+    ret_next = []
     for idx in range(1, len(breaks)):
         label, br = breaks[idx]
         last_label, last_br = breaks[idx - 1]
         if last_label == 'n':
-            ret.append([])
+            if ret_next:
+                ret.append(ret_next)
+                ret_next = []
         if last_label == 'o':
             start = 0
         else:
@@ -88,23 +91,28 @@ def _intersect_1d(breaks):
         last_end = end
         if br == last_br:
             continue
-        ret[-1].append((old_idx, slice(start, end)))
+        ret_next.append((old_idx, slice(start, end)))
         if label == 'o':
             old_idx += 1
             start = 0
-    return tuple(map(tuple, filter(None, ret)))
+
+    if ret_next:
+        ret.append(ret_next)
+
+    return ret
 
 
 def intersect_chunks(old_chunks, new_chunks):
     """
     Make dask.array slices as intersection of old and new chunks.
 
-    >>> intersect_chunks(((4, 4), (2,)),
-    ...                  ((8,), (1, 1)))  # doctest: +NORMALIZE_WHITESPACE
-    ((((0, slice(0, 4, None)), (0, slice(0, 1, None))),
+    >>> intersections = intersect_chunks(((4, 4), (2,)),
+    ...                                  ((8,), (1, 1)))
+    >>> list(intersections)  # doctest: +NORMALIZE_WHITESPACE
+    [(((0, slice(0, 4, None)), (0, slice(0, 1, None))),
       ((1, slice(0, 4, None)), (0, slice(0, 1, None)))),
      (((0, slice(0, 4, None)), (0, slice(1, 2, None))),
-      ((1, slice(0, 4, None)), (0, slice(1, 2, None)))))
+      ((1, slice(0, 4, None)), (0, slice(1, 2, None))))]
 
     Parameters
     ----------
@@ -120,10 +128,10 @@ def intersect_chunks(old_chunks, new_chunks):
     sums2 = [sum(n) for n in old_chunks]
     if not sums == sums2:
         raise ValueError('Cannot change dimensions from to %r' % sums2)
-    old_to_new = tuple(
-        _intersect_1d(_breakpoints(cm[0], cm[1])) for cm in zip(cmo, cmn))
-    cross1 = tuple(product(*old_to_new))
-    cross = tuple(chain(tuple(product(*cr)) for cr in cross1))
+    old_to_new = [_intersect_1d(_breakpoints(cm[0], cm[1]))
+                  for cm in zip(cmo, cmn)]
+    cross1 = product(*old_to_new)
+    cross = chain(tuple(product(*cr)) for cr in cross1)
     return cross
 
 
@@ -475,33 +483,44 @@ def _compute_rechunk(x, chunks):
     x2 = dict()
     intermediates = dict()
     token = tokenize(x, chunks)
-    temp_name = 'rechunk-merge-' + token
-    new_index = tuple(product(*(tuple(range(len(n))) for n in chunks)))
-    for flat_idx, cross1 in enumerate(crossed):
-        new_idx = new_index[flat_idx]
-        key = (temp_name,) + new_idx
-        cr2 = iter(cross1)
-        old_blocks = [[ind for ind, _ in cr] for cr in cross1]
-        subdims = [len(set([ss[i] for ss in old_blocks])) for i in range(ndim)]
-        rec_cat_arg = np.empty(subdims).tolist()
-        inds_in_block = product(*[range(s) for s in subdims])
-        for old_block in old_blocks:
-            ind_slics = next(cr2)
-            old_inds = [[s[0] for s in ind_slics] for i in range(ndim)]
-            # list of nd slices
-            slic = [[s[1] for s in ind_slics] for i in range(ndim)]
-            ind_in_blk = next(inds_in_block)
-            temp = rec_cat_arg
-            for i in range(ndim - 1):
-                temp = getitem(temp, ind_in_blk[i])
-            for ind, slc in zip(old_inds, slic):
-                name = (('rechunk-split-' + token, ) + tuple(ind) +
-                        sum([(s.start, s.stop) for s in slc], ()))
-                intermediates[name] = (getitem, (x.name,) + tuple(ind), tuple(slc))
-                temp[ind_in_blk[-1]] = name
-        x2[key] = (concatenate3, rec_cat_arg)
+    merge_temp_name = 'rechunk-merge-' + token
+    split_temp_name = 'rechunk-split-' + token
+    split_name_suffixes = count()
+
+    # Pre-allocate old block references, to allow re-use and reduce the
+    # graph's memory footprint a bit.
+    old_blocks = np.empty([len(c) for c in x.chunks], dtype='O')
+    for index in np.ndindex(old_blocks.shape):
+        old_blocks[index] = (x.name,) + index
+
+    # Iterate over all new blocks
+    new_index = product(*(range(len(c)) for c in chunks))
+
+    for new_idx, cross1 in zip(new_index, crossed):
+        key = (merge_temp_name,) + new_idx
+        old_block_indices = [[cr[i][0] for cr in cross1] for i in range(ndim)]
+        subdims1 = [len(set(old_block_indices[i]))
+                    for i in range(ndim)]
+
+        rec_cat_arg = np.empty(subdims1, dtype='O')
+        rec_cat_arg_flat = rec_cat_arg.flat
+
+        # Iterate over the old blocks required to build the new block
+        for rec_cat_index, ind_slices in enumerate(cross1):
+            old_block_index, slices = zip(*ind_slices)
+            name = (split_temp_name, next(split_name_suffixes))
+            intermediates[name] = (getitem, old_blocks[old_block_index], slices)
+            rec_cat_arg_flat[rec_cat_index] = name
+
+        assert rec_cat_index == rec_cat_arg.size - 1
+        # New block is formed by concatenation of sliced old blocks
+        x2[key] = (concatenate3, rec_cat_arg.tolist())
+
+    assert new_idx == tuple(len(c) - 1 for c in chunks)
+    del old_blocks, new_index
+
     x2 = merge(x.dask, x2, intermediates)
-    return Array(x2, temp_name, chunks, dtype=x.dtype)
+    return Array(x2, merge_temp_name, chunks, dtype=x.dtype)
 
 
 class _PrettyBlocks(object):


Test patch:
diff --git a/dask/array/tests/test_rechunk.py b/dask/array/tests/test_rechunk.py
--- a/dask/array/tests/test_rechunk.py
+++ b/dask/array/tests/test_rechunk.py
@@ -21,17 +21,17 @@ def test_rechunk_internals_1():
                ('o', 3), ('o', 4), ('o', 5), ('n', 6), ('n', 7))
     assert breaks[1] == answer2
     i1d = [_intersect_1d(b) for b in breaks]
-    answer3 = (((0, slice(0, 1, None)), ),
-               ((0, slice(1, 2, None)), ),
-               ((0, slice(2, 4, None)), ))
+    answer3 = [[(0, slice(0, 1))],
+               [(0, slice(1, 2))],
+               [(0, slice(2, 4))]]
     assert i1d[0] == answer3
-    answer4 = (((0, slice(0, 1, None)), ),
-               ((1, slice(0, 1, None)),
-                (2, slice(0, 1, None)),
-                (3, slice(0, 1, None)),
-                (4, slice(0, 1, None)),
-                (5, slice(0, 1, None))),
-               ((5, slice(1, 2, None)), ))
+    answer4 = [[(0, slice(0, 1))],
+               [(1, slice(0, 1)),
+                (2, slice(0, 1)),
+                (3, slice(0, 1)),
+                (4, slice(0, 1)),
+                (5, slice(0, 1))],
+               [(5, slice(1, 2))]]
     assert i1d[1] == answer4
 
 
@@ -39,12 +39,13 @@ def test_intersect_1():
     """ Convert 1 D chunks"""
     old = ((10, 10, 10, 10, 10), )
     new = ((25, 5, 20), )
-    answer = ((((0, slice(0, 10, None)), ),
-              ((1, slice(0, 10, None)), ),
-              ((2, slice(0, 5, None)), )),
-              (((2, slice(5, 10, None)), ), ),
-              (((3, slice(0, 10, None)), ), ((4, slice(0, 10, None)), )))
-    cross = intersect_chunks(old_chunks=old, new_chunks=new)
+    answer = [(((0, slice(0, 10)), ),
+              ((1, slice(0, 10)), ),
+              ((2, slice(0, 5)), )),
+              (((2, slice(5, 10)), ), ),
+              (((3, slice(0, 10)), ), ((4, slice(0, 10)), ))
+              ]
+    cross = list(intersect_chunks(old_chunks=old, new_chunks=new))
     assert answer == cross
 
 
@@ -52,13 +53,14 @@ def test_intersect_2():
     """ Convert 1 D chunks"""
     old = ((20, 20, 20, 20, 20), )
     new = ((58, 4, 20, 18),)
-    answer = ((((0, slice(0, 20, None)), ),
-              ((1, slice(0, 20, None)), ),
-              ((2, slice(0, 18, None)), )),
-              (((2, slice(18, 20, None)), ), ((3, slice(0, 2, None)), )),
-              (((3, slice(2, 20, None)), ), ((4, slice(0, 2, None)), )),
-              (((4, slice(2, 20, None)), ), ))
-    cross = intersect_chunks(old_chunks=old, new_chunks=new)
+    answer = [(((0, slice(0, 20)), ),
+              ((1, slice(0, 20)), ),
+              ((2, slice(0, 18)), )),
+              (((2, slice(18, 20)), ), ((3, slice(0, 2)), )),
+              (((3, slice(2, 20)), ), ((4, slice(0, 2)), )),
+              (((4, slice(2, 20)), ), )
+              ]
+    cross = list(intersect_chunks(old_chunks=old, new_chunks=new))
     assert answer == cross