dask-image imread v0.5.0 not working with dask distributed Client & napari

@kpasko made a bug report https://github.com/napari/napari/issues/2304 but it turns out this is a problem caused by dask-image. I’ve copied the contents of the report into this issue (sadly I’m unable to transfer issues between different organisations).

What happened:

TypeError: can not serialize ‘function’ object

In distributed/client.py line 2635 futures = self._graph_to_futures

What you expected to happen:

successful image viewing

Minimal Complete Verifiable Example:

(Edited)

from dask.distributed import Client
import napari
from dask_image.imread import imread

client = Client()
data = imread('./*.tif')
napari.view_image(data)

Anything else we need to know?: Works fine when not initializing client, i.e.

from dask.distributed import Client
import napari
from dask_image.imread import imread

data = imread('./*.tif')
napari.view_image(data)

works as expected

Environment:

• Napari/Dask version: dask 2021.2.0 pyhd8ed1ab_0 conda-forge dask-core 2021.2.0 pyhd8ed1ab_0 conda-forge dask-image 0.5.0 pyh44b312d_0 conda-forge distributed 2021.2.0 py39h6e9494a_0 conda-forge napari 0.4.5 pyhd8ed1ab_0 conda-forge napari-console 0.0.3 pyhd8ed1ab_0 conda-forge napari-plugin-engine 0.1.9 py39h6e9494a_1 conda-forge napari-svg 0.1.4 py_0 conda-forge

• Python version: python 3.9.2 h2502468_0_cpython conda-forge

• Operating System: OS X 11.2.1

• Install method (conda, pip, source): conda

[m-albert]

@jakirkham @jrbourbeau I can confirm that https://github.com/dask/dask/pull/7353 fixes all problems discussed in this issue and in particular the original dask-image-napari workflow 🎉

[m-albert]

Okay I narrowed it down to (without napari):

import dask.array as da
import numpy as np

from dask.distributed import Client
import napari
client = Client()

xn = np.random.randint(0, 100, (2, 4, 4))
xd = da.from_array(xn, chunks=(1, 2, 2))

# fails
# def func(block_info=None):
#     return np.random.randint(0, 100, (1, 2, 2))

# works
def func():
    return np.random.randint(0, 100, (1, 2, 2))

xm = da.map_blocks(func, chunks=xd.chunks, dtype=xd.dtype)

from dask.core import flatten
keyset = set(flatten(xm.__dask_keys__()))
xm.dask.__dask_distributed_pack__(client, keyset)

Same error:

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-93-163cc014c931> in <module>
     21 from dask.core import flatten
     22 keyset = set(flatten(xm.__dask_keys__()))
---> 23 xm.dask.__dask_distributed_pack__(client, keyset)

~/miniconda3/envs/dask_image_delme/lib/python3.9/site-packages/dask/highlevelgraph.py in __dask_distributed_pack__(self, client, client_keys)
    942                 }
    943             )
--> 944         return dumps_msgpack({"layers": layers})
    945 
    946     @staticmethod

~/miniconda3/envs/dask_image_delme/lib/python3.9/site-packages/distributed/protocol/core.py in dumps_msgpack(msg, compression)
    161     """
    162     header = {}
--> 163     payload = msgpack.dumps(msg, default=msgpack_encode_default, use_bin_type=True)
    164 
    165     fmt, payload = maybe_compress(payload, compression=compression)

~/miniconda3/envs/dask_image_delme/lib/python3.9/site-packages/msgpack/__init__.py in packb(o, **kwargs)
     33     See :class:`Packer` for options.
     34     """
---> 35     return Packer(**kwargs).pack(o)
     36 
     37 

msgpack/_packer.pyx in msgpack._cmsgpack.Packer.pack()

msgpack/_packer.pyx in msgpack._cmsgpack.Packer.pack()

msgpack/_packer.pyx in msgpack._cmsgpack.Packer.pack()

msgpack/_packer.pyx in msgpack._cmsgpack.Packer._pack()

msgpack/_packer.pyx in msgpack._cmsgpack.Packer._pack()

msgpack/_packer.pyx in msgpack._cmsgpack.Packer._pack()

msgpack/_packer.pyx in msgpack._cmsgpack.Packer._pack()

msgpack/_packer.pyx in msgpack._cmsgpack.Packer._pack()

msgpack/_packer.pyx in msgpack._cmsgpack.Packer._pack()

TypeError: can not serialize 'function' object