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:
cc @rjzamora

On Tue, Mar 23, 2021 at 3:29 AM JSKenyon ***@***.***> wrote:

> *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*:
>
> import dask.array as dafrom 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
>
> —
> 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/7449>, or unsubscribe
> <https://github.com/notifications/unsubscribe-auth/AACKZTASJWAVN5RLPBYNBWLTFBGPZANCNFSM4ZUWHDNQ>
> .
>

Thanks for raising @JSKenyon - I'll check it out
Hmm - This issues is exposing a general HLG/Blockwise challenge related to complex/nested (i.e. inlined) task serialization...

Since the primary callable object is already serialized by the time it reaches the scheduler, in `make_blockwise_graph`, we are replacing the usual `dumps_task` call with [an explicit `{"function": , "args": }` dictionary](https://github.com/dask/dask/blob/f6e9e81903bd2416125a97cc3856c63bb792bddf/dask/blockwise.py#L927).   In cases without task inlining, this is fine, because it is exactly what [dumps_task](https://github.com/dask/distributed/blob/98148cbfee818889f683c4be0e591d33a59ffdf8/distributed/worker.py#L3515) would do anyway.  When there **is** a "complex" task, however, `dumps_task` uses `to_serialize` instead (and the Blockwise behavior is just wrong).

@jrbourbeau @jakirkham - I'll look into this, but do you have any high-level intuition? Might it be possible to use `to_serialize` in this case, even if the primary function is already serialized?

Resolving patch:
diff --git a/dask/blockwise.py b/dask/blockwise.py
--- a/dask/blockwise.py
+++ b/dask/blockwise.py
@@ -342,6 +342,7 @@ def is_materialized(self):
     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 @@ def __dask_distributed_pack__(
         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
--- a/dask/tests/test_distributed.py
+++ b/dask/tests/test_distributed.py
@@ -379,6 +379,33 @@ def add(x, y, z, extra_arg):
     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")