Repository: dask/dask

Issue:
Blockwise with `concatenate=True` broken with LocalCluster + `optimize_graph=False` 
**What happened**:
When using a LocalCluster (and probably the distributed scheduler in general), setting `optimize_graph=False` causes strange (and incorrect) behaviour in Blockwise with `concatentate=True`. 

**What you expected to happen**:
Blockwise with `concatenate=True` should function as normal.

**Minimal Complete Verifiable Example**:

```python
import dask.array as da
from distributed import LocalCluster, Client


def f(x, y):

    print(y)

    return x


if __name__ == '__main__':

    cluster = LocalCluster(processes=False,
                           n_workers=1,
                           threads_per_worker=1)

    client = Client(cluster)

    x = da.array([0])
    y = da.array([[0]])

    z = da.blockwise(
        f, ("i"),
        x, ("i"),
        y, ("ij"),
        dtype=x.dtype,
        concatenate=True)

    da.compute(z, optimize_graph=False)
```

**Anything else we need to know?**:
The example above will print something like:
```
(<function concatenate_axes at 0x7f153d45c310>, [array([[0]])], [1])
```
This is not what is expected - we expect an array. Note that the above example will work as normal when:
- `optimize_graph=True`
- A different scheduler is used.

Apologies if this should have been posted on `distributed` instead.

**Environment**:

- Dask version: 2021.3.0
- Distributed version: 2021.3.0
- Python version: 3.8.8
- Operating System: Pop!_OS 18.04 LTS
- Install method (conda, pip, source): pip


Discussion:
jrbourbeau: Just pushed a commit to temporarily point CI to https://github.com/dask/distributed/pull/4641 in order to illustrate that tests pass here when we include the `distributed` changes -- hope that's okay @rjzamora 
rjzamora: @jrbourbeau - What is the best way to deal with the distributed min-dep changing to the next release (2021.04.1)?  For now, I just changed the min-dep environment to install distributed from source.
rjzamora: @jrbourbeau - I am okay with merging #7415 very soon after the next release, but it would be nice to get this fix in as soon as possible.  Do you expect there to be problems with this plan?
jrbourbeau: That sounds good to me. The issue here is we're getting consistent failures in the `mindeps-distributed` build due to calling `pytest.importorskip` in `gen_cluster` (we're running into this in other places like https://github.com/dask/dask/pull/7524 too).

I tried to determine what the underlying issue is yesterday but wasn't able to figure it out yet. I did observe that the issue is only present on Python 3.7 so I tried to temporarily bump the Python version to used in the `mindeps-distributed` build to 3.8 over in [this branch](https://github.com/jrbourbeau/dask/tree/min-dep-dist-py38) but ran into some conda solve conflicts. That's where things currently stand. I'll try to pick this up again today (or you should feel free to if you have the bandwidth). @jsignell also suggested we might try moving tests into separate modules and skipping entire modules instead of individual tests. I suspect this would also work as a quickfix, but long-term I'd prefer to keep the current file structure. 
jrbourbeau: Alright, I ended up changing the location of the test added here so it's no longer the last test in this module. That should hopefully result in CI passing 🤞It's unfortunate that test ordering currently matters, but the actual change here is small and benign
rjzamora: Interesting - Would it help for us to add a module-check decorator (so that we skip outside the `gen_cluster`)?
jrbourbeau: Hmm good idea, something like that would also work. Though from experimenting locally it looks like `find_spec` returns a non-`None` output for `dask.array` even when I don't have `numpy` installed. I'm fine trying to include a module check decorator or just keeping things as they currently are in this PR 
jrbourbeau: Just pushed a commit to resolve a merge conflict. Planning to merge this PR after CI finishes 

Patch:
diff --git a/continuous_integration/environment-mindeps-distributed.yaml b/continuous_integration/environment-mindeps-distributed.yaml
index c41e78995..8742c65f6 100644
--- a/continuous_integration/environment-mindeps-distributed.yaml
+++ b/continuous_integration/environment-mindeps-distributed.yaml
@@ -10,7 +10,10 @@ dependencies:
   - fsspec=0.6.0
   - toolz=0.8.2
   # optional dependencies pulled in by pip install dask[distributed]
-  - distributed=2021.03.0
+  # TODO: Use distributed-2021.04.1 when it becomes available
+  - pip
+  - pip:
+    - git+https://github.com/dask/distributed.git
   # test dependencies
   - pytest
   - pytest-rerunfailures
diff --git a/dask/blockwise.py b/dask/blockwise.py
index a5138f38b..443ad6d68 100644
--- a/dask/blockwise.py
+++ b/dask/blockwise.py
@@ -342,6 +342,7 @@ class Blockwise(Layer):
     def __dask_distributed_pack__(
         self, all_hlg_keys, known_key_dependencies, client, client_keys
     ):
+        from distributed.protocol import to_serialize
         from distributed.protocol.serialize import import_allowed_module
         from distributed.utils import CancelledError
         from distributed.utils_comm import unpack_remotedata
@@ -363,7 +364,9 @@ class Blockwise(Layer):
         dsk = (SubgraphCallable(dsk, self.output, tuple(keys2)),)
         dsk, dsk_unpacked_futures = unpack_remotedata(dsk, byte_keys=True)
 
-        func = dumps_function(dsk[0])
+        # Dump the function if concatenate is False, because
+        # we will not need to construct a nested task
+        func = to_serialize(dsk[0]) if self.concatenate else dumps_function(dsk[0])
         func_future_args = dsk[1:]
 
         indices = list(toolz.concat(indices2))
@@ -819,8 +822,8 @@ def make_blockwise_graph(
         key_deps = {}
 
     if deserializing:
-        from distributed.protocol.serialize import import_allowed_module
-        from distributed.worker import dumps_function, warn_dumps
+        from distributed.protocol.serialize import import_allowed_module, to_serialize
+        from distributed.worker import dumps_function
     else:
         from importlib import import_module as import_allowed_module
 
@@ -915,21 +918,30 @@ def make_blockwise_graph(
         if deserializing:
             deps.update(func_future_args)
             args += list(func_future_args)
+
+        if deserializing and not concatenate:
+            # Construct a function/args/kwargs dict if we
+            # do not have a nested task (i.e. concatenate=False).
+            # TODO: Avoid using the iterate_collection-version
+            # of to_serialize if we know that are no embeded
+            # Serialized/Serialize objects in args and/or kwargs.
             if kwargs:
-                val = {
+                dsk[out_key] = {
                     "function": dumps_function(apply),
-                    "args": warn_dumps(args),
-                    "kwargs": warn_dumps(kwargs2),
+                    "args": to_serialize(args),
+                    "kwargs": to_serialize(kwargs2),
                 }
             else:
-                val = {"function": func, "args": warn_dumps(args)}
+                dsk[out_key] = {"function": func, "args": to_serialize(args)}
         else:
             if kwargs:
                 val = (apply, func, args, kwargs2)
             else:
                 args.insert(0, func)
                 val = tuple(args)
-        dsk[out_key] = val
+            # May still need to serialize (if concatenate=True)
+            dsk[out_key] = to_serialize(val) if deserializing else val
+
         if return_key_deps:
             key_deps[out_key] = deps
 


Test patch:
diff --git a/dask/tests/test_distributed.py b/dask/tests/test_distributed.py
index 426962c98..2e520872b 100644
--- a/dask/tests/test_distributed.py
+++ b/dask/tests/test_distributed.py
@@ -379,6 +379,33 @@ async def test_combo_of_layer_types(c, s, a, b):
     assert res == 21
 
 
+@gen_cluster(client=True)
+async def test_blockwise_concatenate(c, s, a, b):
+    """Test a blockwise operation with concatenated axes"""
+    da = pytest.importorskip("dask.array")
+    np = pytest.importorskip("numpy")
+
+    def f(x, y):
+        da.assert_eq(y, [[0, 1, 2]])
+        return x
+
+    x = da.from_array(np.array([0, 1, 2]))
+    y = da.from_array(np.array([[0, 1, 2]]))
+    z = da.blockwise(
+        f,
+        ("i"),
+        x,
+        ("i"),
+        y,
+        ("ij"),
+        dtype=x.dtype,
+        concatenate=True,
+    )
+
+    await c.compute(z, optimize_graph=False)
+    da.assert_eq(z, x)
+
+
 @gen_cluster(client=True)
 async def test_map_partitions_partition_info(c, s, a, b):
     dd = pytest.importorskip("dask.dataframe")