Repository: dask/dask

Issue:
da.map_blocks with drop_axis stopped working in dask 1.1.4 
Hi!

First and foremost a big Thank You for the amazing work that you're doing with dask! It's a fantastic piece of software! 

We're using dask as parallelization engine for a large-scale data analysis package (still in early development). Our computational routines work on independent chunks of one big source data-set and then write to a single file (accessible to all workers). We used `map_blocks` for this in dask 1.0.0 and things worked great but unfortunately the same code does not work any more in dask 1.1.4. Here's a MCVE:

```python
import dask
import dask.array as da
import numpy as np

# in reality used to write `chk` to a (pre-formatted) file
def da_dummy(chk, block_info=None):
    idx = block_info[0]["chunk-location"][-1]
    #
    # use `idx` to store `chk` at correct location
    # 
    return idx

np_a = np.ones((10,50*20))
a = da.from_array(np_a, chunks=(10,20))

# works in dask 1.0.0, raises an IndexError in dask 1.1.4
result = a.map_blocks(da_dummy, dtype="int", chunks=(1,), drop_axis=[0])
res = result.compute()
```
In dask 1.0.0 this snippet produces an array `res` that looks like this:
```python
array([ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11, 12, 13, 14, 15, 16,
       17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33,
       34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49])
```
In dask 1.1.4 the same code triggers an `IndexError: index 2 is out of bounds for axis 0 with size 2`, here the full traceback:

<details>

```python
---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
~/Documents/job/dask_fail.py in <module>
     15 
     16 # works in dask 1.0.0, raises an IndexError in dask 1.1.4
---> 17 result = a.map_blocks(da_dummy, dtype="int", chunks=(1,), drop_axis=[0])
     18 res = result.compute()
     19 

/usr/local/software/anaconda/envs/p3/lib/python3.6/site-packages/dask/array/core.py in map_blocks(self, func, *args, **kwargs)
   1569     @wraps(map_blocks)
   1570     def map_blocks(self, func, *args, **kwargs):
-> 1571         return map_blocks(func, self, *args, **kwargs)
   1572 
   1573     def map_overlap(self, func, depth, boundary=None, trim=True, **kwargs):

/usr/local/software/anaconda/envs/p3/lib/python3.6/site-packages/dask/array/core.py in map_blocks(func, *args, **kwargs)
    534                                            for ij, j in enumerate(old_k[1:])],
    535                         'chunk-location': old_k[1:]}
--> 536                     for i in shapes}
    537 
    538             v = copy.copy(v)  # Need to copy and unpack subgraph callable

/usr/local/software/anaconda/envs/p3/lib/python3.6/site-packages/dask/array/core.py in <dictcomp>(.0)
    534                                            for ij, j in enumerate(old_k[1:])],
    535                         'chunk-location': old_k[1:]}
--> 536                     for i in shapes}
    537 
    538             v = copy.copy(v)  # Need to copy and unpack subgraph callable

/usr/local/software/anaconda/envs/p3/lib/python3.6/site-packages/dask/array/core.py in <listcomp>(.0)
    532                         'num-chunks': num_chunks[i],
    533                         'array-location': [(starts[i][ij][j], starts[i][ij][j + 1])
--> 534                                            for ij, j in enumerate(old_k[1:])],
    535                         'chunk-location': old_k[1:]}
    536                     for i in shapes}

IndexError: index 2 is out of bounds for axis 0 with size 2
```

</details>

A (desperate) attempt at somehow working around this was to not use `drop_axis` and instead leave the chunk-size as is and return a bogus tuple, but no luck either:

```python

import dask
import dask.array as da
import numpy as np

np_a = np.ones((10,50*20))
a = da.from_array(np_a, chunks=(10,20))

# desperate workaround for dask 1.1.4 - still not working
def da_dummy2(chk, block_info=None):
    print(block_info) # show what's going on in here...
    idx = block_info[0]["chunk-location"][-1]
    return (idx, 1)
result = a.map_blocks(da_dummy2, dtype="int", chunks=(1, 1))
res = result.compute()
```

This gives an `IndexError: tuple index out of range`, full traceback here:

<details>

```python
---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
~/Documents/job/dask_fail.py in <module>
     20     return (idx, 1)
     21 result = a.map_blocks(da_dummy2, dtype="int", chunks=(1, 1))
---> 22 res = result.compute()
     23 

/usr/local/software/anaconda/envs/p3/lib/python3.6/site-packages/dask/base.py in compute(self, **kwargs)
    154         dask.base.compute
    155         """
--> 156         (result,) = compute(self, traverse=False, **kwargs)
    157         return result
    158 

/usr/local/software/anaconda/envs/p3/lib/python3.6/site-packages/dask/base.py in compute(*args, **kwargs)
    397     postcomputes = [x.__dask_postcompute__() for x in collections]
    398     results = schedule(dsk, keys, **kwargs)
--> 399     return repack([f(r, *a) for r, (f, a) in zip(results, postcomputes)])
    400 
    401 

/usr/local/software/anaconda/envs/p3/lib/python3.6/site-packages/dask/base.py in <listcomp>(.0)
    397     postcomputes = [x.__dask_postcompute__() for x in collections]
    398     results = schedule(dsk, keys, **kwargs)
--> 399     return repack([f(r, *a) for r, (f, a) in zip(results, postcomputes)])
    400 
    401 

/usr/local/software/anaconda/envs/p3/lib/python3.6/site-packages/dask/array/core.py in finalize(results)
    777     while isinstance(results2, (tuple, list)):
    778         if len(results2) > 1:
--> 779             return concatenate3(results)
    780         else:
    781             results2 = results2[0]

/usr/local/software/anaconda/envs/p3/lib/python3.6/site-packages/dask/array/core.py in concatenate3(arrays)
   3495     if not ndim:
   3496         return arrays
-> 3497     chunks = chunks_from_arrays(arrays)
   3498     shape = tuple(map(sum, chunks))
   3499 

/usr/local/software/anaconda/envs/p3/lib/python3.6/site-packages/dask/array/core.py in chunks_from_arrays(arrays)
   3325 
   3326     while isinstance(arrays, (list, tuple)):
-> 3327         result.append(tuple([shape(deepfirst(a))[dim] for a in arrays]))
   3328         arrays = arrays[0]
   3329         dim += 1

/usr/local/software/anaconda/envs/p3/lib/python3.6/site-packages/dask/array/core.py in <listcomp>(.0)
   3325 
   3326     while isinstance(arrays, (list, tuple)):
-> 3327         result.append(tuple([shape(deepfirst(a))[dim] for a in arrays]))
   3328         arrays = arrays[0]
   3329         dim += 1

IndexError: tuple index out of range
```

</details>

Any help with this is greatly appreciated. Thank you! 

Discussion:
bmerry: @jakirkham would you be able to review this?
TomAugspurger: Thank @bmerry, this looks good.

> This may have a minor performance impact because there is
now an extra per-chunk dictionary being thrown around.

As a sanity check, things look OK

```python
import dask.array as da

arr = da.random.random_integers(0, 10, size=(100_000, 100), chunks=100)

%timeit arr.map_blocks(lambda x, block_info=None: x)
%timeit arr.map_blocks(lambda x, block_info=None: x[:, 0], drop_axis=1)
```

I don't see any slowdown from master. Can you confirm that those should basically be the worst-case scenario for performance, since those functions basically do nothing?
bmerry: The performance drop I was referring to is comparing before and after this patch (both without `drop_axis`), rather than with vs without `drop_axis`. Modifying your test slightly to use more chunks:

```python
import dask.array as da

arr = da.random.random_integers(0, 10, size=(100_000, 100), chunks=10)

%timeit arr.map_blocks(lambda x, block_info=None: x)
```

I get 6.13 s ± 47.1 ms from 1.2.2 and 6.35 s ± 33 ms from this branch. So the difference is statistically significant but not very large.

Patch:
diff --git a/dask/array/core.py b/dask/array/core.py
index f9fa56a15..40c6bc0a6 100644
--- a/dask/array/core.py
+++ b/dask/array/core.py
@@ -536,7 +536,7 @@ def map_blocks(func, *args, **kwargs):
         chunks2 = []
         for i, (c, nb) in enumerate(zip(chunks, out.numblocks)):
             if isinstance(c, tuple):
-                # We only check cases where numblocks > 1. Becuase of
+                # We only check cases where numblocks > 1. Because of
                 # broadcasting, we can't (easily) validate the chunks
                 # when the number of blocks is 1.
                 # See https://github.com/dask/dask/issues/4299 for more.
@@ -555,16 +555,19 @@ def map_blocks(func, *args, **kwargs):
         num_chunks = {}
         shapes = {}
 
-        for i, arg in enumerate(args):
-            if isinstance(arg, Array):
-                starts[i] = [np.cumsum((0,) + c) for c in arg.chunks]
+        for i, (arg, in_ind) in enumerate(argpairs):
+            if in_ind is not None:
                 shapes[i] = arg.shape
-                num_chunks[i] = arg.numblocks
-        for k, v in kwargs.items():
-            if isinstance(v, Array):
-                starts[k] = [np.cumsum((0,) + c) for c in v.chunks]
-                shapes[k] = arg.shape
-                num_chunks[i] = arg.numblocks
+                if drop_axis:
+                    # We concatenate along dropped axes, so we need to treat them
+                    # as if there is only a single chunk.
+                    starts[i] = [(np.cumsum((0,) + arg.chunks[j])
+                                  if ind in out_ind else np.array([0, arg.shape[j]]))
+                                 for j, ind in enumerate(in_ind)]
+                    num_chunks[i] = tuple(len(s) - 1 for s in starts[i])
+                else:
+                    starts[i] = [np.cumsum((0,) + c) for c in arg.chunks]
+                    num_chunks[i] = arg.numblocks
         out_starts = [np.cumsum((0,) + c) for c in out.chunks]
 
         for k, v in dsk.items():
@@ -572,21 +575,16 @@ def map_blocks(func, *args, **kwargs):
             v = v[0]
             [(key, task)] = v.dsk.items()  # unpack subgraph callable
 
-            if new_axis:
-                # anything using the keys in dsk is incorrect, as the
-                # original array doesn't have values for `new_axis`.
-                old_k = tuple(x for i, x in enumerate(k[1:]) if i not in new_axis)
-            else:
-                old_k = k[1:]
-
+            # Get position of chunk, indexed by axis labels
+            location = {out_ind[i]: loc for i, loc in enumerate(k[1:])}
             info = {}
             for i, shape in shapes.items():
                 # Compute chunk key in the array, taking broadcasting into
                 # account. We don't directly know which dimensions are
                 # broadcast, but any dimension with only one chunk can be
                 # treated as broadcast.
-                arr_k = old_k[-len(shape):]
-                arr_k = tuple(j if num_chunks[i][ij] > 1 else 0 for ij, j in enumerate(arr_k))
+                arr_k = tuple(location.get(ind, 0) if num_chunks[i][j] > 1 else 0
+                              for j, ind in enumerate(argpairs[i][1]))
                 info[i] = {'shape': shape,
                            'num-chunks': num_chunks[i],
                            'array-location': [(starts[i][ij][j], starts[i][ij][j + 1])
diff --git a/docs/source/setup.rst b/docs/source/setup.rst
index 60d71f3e3..818a46594 100644
--- a/docs/source/setup.rst
+++ b/docs/source/setup.rst
@@ -84,3 +84,4 @@ how to set up Dask on a variety of local and distributed hardware.
    setup/cloud.rst
    setup/adaptive.rst
    setup/docker.rst
+   setup/custom-startup.rst
diff --git a/docs/source/setup/custom-startup.rst b/docs/source/setup/custom-startup.rst
new file mode 100644
index 000000000..eebe92efb
--- /dev/null
+++ b/docs/source/setup/custom-startup.rst
@@ -0,0 +1,73 @@
+Custom Initialization
+=====================
+
+Often we want to run custom code when we start up or tear down a scheduler or
+worker.  We might do this manually with functions like ``Client.run`` or
+``Client.run_on_scheduler``, but this is error prone and difficult to automate.
+
+To resolve this, Dask includes a few mechanisms to run arbitrary code around
+the lifecycle of a Scheduler or Worker.
+
+Preload Scripts
+---------------
+
+Both ``dask-scheduler`` and ``dask-worker`` support a ``--preload`` option that
+allows custom initialization of each scheduler/worker respectively. A module or
+Python file passed as a ``--preload`` value is guaranteed to be imported before
+establishing any connection. A ``dask_setup(service)`` function is called if
+found, with a ``Scheduler`` or ``Worker`` instance as the argument. As the
+service stops, ``dask_teardown(service)`` is called if present.
+
+To support additional configuration, a single ``--preload`` module may register
+additional command-line arguments by exposing ``dask_setup`` as a  Click_
+command.  This command will be used to parse additional arguments provided to
+``dask-worker`` or ``dask-scheduler`` and will be called before service
+initialization.
+
+.. _Click: http://click.pocoo.org/
+
+
+As an example, consider the following file that creates a
+:doc:`scheduler plugin <plugins>` and registers it with the scheduler
+
+.. code-block:: python
+
+   # scheduler-setup.py
+   import click
+
+   from distributed.diagnostics.plugin import SchedulerPlugin
+
+   class MyPlugin(SchedulerPlugin):
+       def __init__(self, print_count):
+         self.print_count = print_count
+         SchedulerPlugin.__init__(self)
+
+       def add_worker(self, scheduler=None, worker=None, **kwargs):
+           print("Added a new worker at:", worker)
+           if self.print_count and scheduler is not None:
+               print("Total workers:", len(scheduler.workers))
+
+   @click.command()
+   @click.option("--print-count/--no-print-count", default=False)
+   def dask_setup(scheduler, print_count):
+       plugin = MyPlugin(print_count)
+       scheduler.add_plugin(plugin)
+
+We can then run this preload script by referring to its filename (or module name
+if it is on the path) when we start the scheduler::
+
+   dask-scheduler --preload scheduler-setup.py --print-count
+
+
+Worker Lifecycle Plugins
+------------------------
+
+You can also create a class with setup and teardown methods, and register that
+class with the scheduler to give to every worker.
+
+.. currentmodule:: distributed
+
+.. autosummary::
+   Client.register_worker_plugin
+
+.. automethod:: Client.register_worker_plugin


Test patch:
diff --git a/dask/array/tests/test_array_core.py b/dask/array/tests/test_array_core.py
index b4b5cc30d..2da84ff96 100644
--- a/dask/array/tests/test_array_core.py
+++ b/dask/array/tests/test_array_core.py
@@ -1115,7 +1115,7 @@ def test_map_blocks_block_info():
     assert_eq(z, x + x + 1 + 100)
 
 
-def test_map_blocks_block_info_with_axis():
+def test_map_blocks_block_info_with_new_axis():
     # https://github.com/dask/dask/issues/4298
     values = da.from_array(np.array(['a', 'a', 'b', 'c']), 2)
 
@@ -1145,6 +1145,40 @@ def test_map_blocks_block_info_with_axis():
     assert_eq(z, np.ones((4, 3), dtype='f8'))
 
 
+def test_map_blocks_block_info_with_drop_axis():
+    # https://github.com/dask/dask/issues/4584
+    values = da.from_array(np.array([[1, 2, 4],
+                                     [8, 16, 32],
+                                     [64, 128, 256],
+                                     [1024, 2048, 4096]], dtype='u4'), (2, 1))
+
+    def func(x, block_info=None):
+        assert set(block_info.keys()) == {0, None}
+        assert block_info[0]['shape'] == (4, 3)
+        # drop_axis concatenates along the dropped dimension, hence not (2, 3)
+        assert block_info[0]['num-chunks'] == (2, 1)
+        assert block_info[None]['shape'] == (4,)
+        assert block_info[None]['num-chunks'] == (2,)
+        assert block_info[None]['chunk-shape'] == (2,)
+        assert block_info[None]['dtype'] == np.dtype('u4')
+
+        assert block_info[0]['chunk-location'] in {(0, 0), (1, 0)}
+
+        if block_info[0]['chunk-location'] == (0, 0):
+            assert block_info[0]['array-location'] == [(0, 2), (0, 3)]
+            assert block_info[None]['chunk-location'] == (0,)
+            assert block_info[None]['array-location'] == [(0, 2)]
+        elif block_info[0]['chunk-location'] == (1, 0):
+            assert block_info[0]['array-location'] == [(2, 4), (0, 3)]
+            assert block_info[None]['chunk-location'] == (1,)
+            assert block_info[None]['array-location'] == [(2, 4)]
+
+        return np.sum(x, axis=1, dtype='u4')
+
+    z = values.map_blocks(func, drop_axis=1, dtype='u4')
+    assert_eq(z, np.array([7, 56, 448, 7168], dtype='u4'))
+
+
 def test_map_blocks_block_info_with_broadcast():
     expected0 = [
         {'shape': (3, 4), 'num-chunks': (1, 2), 'array-location': [(0, 3), (0, 2)], 'chunk-location': (0, 0)},