Repository: dask/dask
Pull request: Handle Blockwise HLG pack/unpack for concatenate=True
One possible solution to the `Blockwise` serialization challenge discussed in [this comment](https://github.com/dask/dask/issues/7449#issuecomment-805073471).  Note that I would certainly prefer a general solution to nested-task serialization.

- [x] Closes #7449
- [x] Tests added / passed
- [x] Passes `black dask` / `flake8 dask`


Problem statement:
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

Hints from issue 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?

Commit message: use simpler SerializedFunction wrapper

Patch to review:
diff --git a/dask/blockwise.py b/dask/blockwise.py
index e116d26c04b..c1bf3588dab 100644
--- a/dask/blockwise.py
+++ b/dask/blockwise.py
@@ -417,6 +417,8 @@ def __dask_distributed_pack__(
 
     @classmethod
     def __dask_distributed_unpack__(cls, state, dsk, dependencies):
+        from .layers import SerializedFunction
+
         # Make sure we convert list items back from tuples in `indices`.
         # The msgpack serialization will have converted lists into
         # tuples, and tuples may be stringified during graph
@@ -427,7 +429,7 @@ def __dask_distributed_unpack__(cls, state, dsk, dependencies):
         ]
 
         layer_dsk, layer_deps = make_blockwise_graph(
-            state["func"],
+            SerializedFunction(state["func"]),
             state["output"],
             state["output_indices"],
             *indices,
@@ -821,8 +823,8 @@ def make_blockwise_graph(
         key_deps = {}
 
     if deserializing:
-        from distributed.worker import warn_dumps, dumps_function
         from distributed.protocol.serialize import import_allowed_module
+        from distributed.protocol import to_serialize
     else:
         from importlib import import_module as import_allowed_module
 
@@ -917,21 +919,12 @@ def make_blockwise_graph(
         if deserializing:
             deps.update(func_future_args)
             args += list(func_future_args)
-            if kwargs:
-                val = {
-                    "function": dumps_function(apply),
-                    "args": warn_dumps(args),
-                    "kwargs": warn_dumps(kwargs2),
-                }
-            else:
-                val = {"function": func, "args": warn_dumps(args)}
+        if kwargs:
+            val = (apply, func, args, kwargs2)
         else:
-            if kwargs:
-                val = (apply, func, args, kwargs2)
-            else:
-                args.insert(0, func)
-                val = tuple(args)
-        dsk[out_key] = val
+            args.insert(0, func)
+            val = tuple(args)
+        dsk[out_key] = to_serialize(val) if deserializing else val
         if return_key_deps:
             key_deps[out_key] = deps
 
diff --git a/dask/layers.py b/dask/layers.py
index 6a55f3e2485..a8efec11b28 100644
--- a/dask/layers.py
+++ b/dask/layers.py
@@ -31,6 +31,16 @@ def __call__(self, *args, **kwargs):
         return import_term(self.function_path)(*args, **kwargs)
 
 
+class SerializedFunction:
+    def __init__(self, func):
+        self.func = func
+
+    def __call__(self, *args, **kwargs):
+        from distributed.worker import loads_function
+
+        return loads_function(self.func)(*args, **kwargs)
+
+
 #
 ##
 ###  Array Layers & Utilities
diff --git a/dask/tests/test_distributed.py b/dask/tests/test_distributed.py
index 4b2bb0ffcb5..f930b2a7fc1 100644
--- a/dask/tests/test_distributed.py
+++ b/dask/tests/test_distributed.py
@@ -398,3 +398,30 @@ async def test_annotation_pack_unpack(c, s, a, b):
     unpacked_hlg = HighLevelGraph.__dask_distributed_unpack__(packed_hlg, annotations)
     annotations = unpacked_hlg["annotations"]
     assert annotations == {"workers": {"n": ("alice",)}}
+
+
+@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)


Code review comments:
diff_hunk: @@ -917,21 +919,12 @@ def make_blockwise_graph(
         if deserializing:
             deps.update(func_future_args)
             args += list(func_future_args)
-            if kwargs:
-                val = {
-                    "function": dumps_function(apply),
-                    "args": warn_dumps(args),
-                    "kwargs": warn_dumps(kwargs2),
-                }
-            else:
-                val = {"function": func, "args": warn_dumps(args)}
+        if kwargs:
+            val = (apply, func, args, kwargs2)
path: dask/blockwise.py
text: @author:
@user1 and @user2-r-rose - I am able to handle in-lined tasks and avoid the extra `dumps` calls here if I wrap the function in something simple (to make it "callable" again).  However, I am thinking that I am still missing the best way to use the existing serialization machinery (otherwise, the best long-term solution may require improvements in `distributed`/`Serialize`).

The current challenge, both here and in #7417, seems to be that we cannot serialize only the function and/or task arguments (or a subset of them) on the client, and then also serialize the complete task on the scheduler (which is necessary to send the task to a worker).   I would expect that the `serialize` logic would come across a `Serialized` object, and incorperate the header/frames into the *new* `Serialized` object so that calling `deserialize` on the final object would also deserialize the embedded object.  However, for some reason, this does not seem to be the case. (note that I'm still looking into this...)

@author:
Ah- It looks like this is another case where we just need to add an `iterate_collection=` argument to `serialize`.  If we require the serialization procedure to treat the elements of a collection separately, the nested serialization seems to work for me :)

diff_hunk: @@ -917,21 +919,12 @@ def make_blockwise_graph(
         if deserializing:
             deps.update(func_future_args)
             args += list(func_future_args)
-            if kwargs:
-                val = {
-                    "function": dumps_function(apply),
-                    "args": warn_dumps(args),
-                    "kwargs": warn_dumps(kwargs2),
-                }
-            else:
-                val = {"function": func, "args": warn_dumps(args)}
+        if kwargs:
+            val = (apply, func, args, kwargs2)
         else:
-            if kwargs:
-                val = (apply, func, args, kwargs2)
-            else:
-                args.insert(0, func)
-                val = tuple(args)
-        dsk[out_key] = val
+            args.insert(0, func)
+            val = tuple(args)
+        dsk[out_key] = to_serialize(val) if deserializing else val
path: dask/blockwise.py
text: @author:
Note that, if [distributed#4641](https://github.com/dask/distributed/pull/4641) gets in, we can do something like:

```
dsk[out_key] = Serialized(*serialize(val, iterate_collection=True)) if deserializing else val
```

and then the `SerializedFunction` wrapper becomes completely unnecessary :)

@user1:
Just a heads up, when using https://github.com/dask/distributed/pull/4575, `to_serialize` should work here:

```python
dsk[out_key] = to_serialize(val) if deserializing else val
```

@author:
Ah, cool - thanks!  I'll give that a try.

@user2:
That PR still has other issues though. So something to be aware of

@author:
>That PR still has other issues though. So something to be aware of

Thanks for the heads up John.  Is there anything specific that we can do to help?  It seems that Blockwise-IO serialization may be blocked in dask/dask until something like your PR gets through in dask/distributed :)