[AstroPy] astropy Quantity, dask serializer conflict?

[Nathaniel J Livesey]

Dear all,

   I’m encountering issues with mixing dask and astropy Quantity.  If I run an astropy Quantity instance through the default dask serializer/deserializer, what emerges is a ndarray.  To make it work I have to disable the “dask” serializer method and have it use only the “pickle” one.

   See the MWE below:
==============
import astropy.units as units
import dask.distributed as distributed
import numpy as np

x = np.arange(5) << units.s
print(f"Input is {type(x)}: {x}")

# In the first case, serializers defaults to ["dask", "pickle"]
header, frames = distributed.protocol.serialize(x)
y = distributed.protocol.deserialize(header, frames)
print(f"Default serialization gives {type(y)}: {y}")

# Here we disable the "dask" serializer
header, frames = distributed.protocol.serialize(x, serializers=["pickle"])
y = distributed.protocol.deserialize(header, frames)
print(f"With dask serialization disabled we get {type(y)}: {y}")
===============

Which results in:

Input is <class 'astropy.units.quantity.Quantity'>: [0. 1. 2. 3. 4.] s
Default serialization gives <class 'numpy.ndarray'>: [0. 1. 2. 3. 4.]
With dask serialization disabled we get <class 'astropy.units.quantity.Quantity'>: [0. 1. 2. 3. 4.] s

===============
I’m currently running astropy 4.3.1, but the same thing happens with the latest version from github.  I’m using python 3.9, dask 2021.10.0 and numpy 1.20.3

Many thanks for any pointers.

Yours,

Nathaniel

P.S. I have a work around that disables the “dask” serializer when I start the dask Client.  However, that sometimes doesn’t work because something else (I suspect xarray) seems to overwrite that setting.

Nathaniel Livesey
Mail Stop 183-701, Jet Propulsion Laboratory
4800 Oak Grove Drive, Pasadena, California 91109.
Phone: +1 818 354-4214   Fax: +1 818 393-5065