astropy: Writing (distributed) Dask array fails

Description

Dask array fails to be written to FITS, at least when using dask.distributed. This doesn’t occur for me if I don’t create a dask.distributed Client. That is, the example given in the 'What’s new in 4.1` works for me.

Expected behavior

As per version 4.1, writing dask arrays is supported.

Actual behavior

The following error occurred:

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
~/miniconda3/lib/python3.7/site-packages/distributed/protocol/pickle.py in dumps(x, buffer_callback, protocol)
     48         buffers.clear()
---> 49         result = pickle.dumps(x, **dump_kwargs)
     50         if len(result) < 1000:

TypeError: can't pickle _thread.lock objects

During handling of the above exception, another exception occurred:

TypeError                                 Traceback (most recent call last)
<ipython-input-44-23f936e3e991> in <module>
      3 from astropy.io import fits
      4 hdu = fits.PrimaryHDU(data=array)
----> 5 hdu.writeto('test_dask.fits', overwrite=True)

~/miniconda3/lib/python3.7/site-packages/astropy/utils/decorators.py in wrapper(*args, **kwargs)
    533                     warnings.warn(message, warning_type, stacklevel=2)
    534 
--> 535             return function(*args, **kwargs)
    536 
    537         return wrapper

~/miniconda3/lib/python3.7/site-packages/astropy/io/fits/hdu/base.py in writeto(self, name, output_verify, overwrite, checksum)
    370         hdulist = HDUList([self])
    371         hdulist.writeto(name, output_verify, overwrite=overwrite,
--> 372                         checksum=checksum)
    373 
    374     @classmethod

~/miniconda3/lib/python3.7/site-packages/astropy/utils/decorators.py in wrapper(*args, **kwargs)
    533                     warnings.warn(message, warning_type, stacklevel=2)
    534 
--> 535             return function(*args, **kwargs)
    536 
    537         return wrapper

~/miniconda3/lib/python3.7/site-packages/astropy/io/fits/hdu/hdulist.py in writeto(self, fileobj, output_verify, overwrite, checksum)
    941             for hdu in self:
    942                 hdu._prewriteto(checksum=checksum)
--> 943                 hdu._writeto(hdulist._file)
    944                 hdu._postwriteto()
    945         hdulist.close(output_verify=output_verify, closed=closed)

~/miniconda3/lib/python3.7/site-packages/astropy/io/fits/hdu/base.py in _writeto(self, fileobj, inplace, copy)
    675 
    676         with _free_space_check(self, dirname):
--> 677             self._writeto_internal(fileobj, inplace, copy)
    678 
    679     def _writeto_internal(self, fileobj, inplace, copy):

~/miniconda3/lib/python3.7/site-packages/astropy/io/fits/hdu/base.py in _writeto_internal(self, fileobj, inplace, copy)
    681         if not inplace or self._new:
    682             header_offset, _ = self._writeheader(fileobj)
--> 683             data_offset, data_size = self._writedata(fileobj)
    684 
    685             # Set the various data location attributes on newly-written HDUs

~/miniconda3/lib/python3.7/site-packages/astropy/io/fits/hdu/base.py in _writedata(self, fileobj)
    613         if self._data_loaded or self._data_needs_rescale:
    614             if self.data is not None:
--> 615                 size += self._writedata_internal(fileobj)
    616             # pad the FITS data block
    617             if size > 0:

~/miniconda3/lib/python3.7/site-packages/astropy/io/fits/hdu/image.py in _writedata_internal(self, fileobj)
    621             return size
    622         elif isinstance(self.data, DaskArray):
--> 623             return self._writeinternal_dask(fileobj)
    624         else:
    625             # Based on the system type, determine the byteorders that

~/miniconda3/lib/python3.7/site-packages/astropy/io/fits/hdu/image.py in _writeinternal_dask(self, fileobj)
    707                                 buffer=outmmap)
    708 
--> 709             output.store(outarr, lock=True, compute=True)
    710         finally:
    711             if should_close:

~/miniconda3/lib/python3.7/site-packages/dask/array/core.py in store(self, target, **kwargs)
   1387     @wraps(store)
   1388     def store(self, target, **kwargs):
-> 1389         r = store([self], [target], **kwargs)
   1390 
   1391         if kwargs.get("return_stored", False):

~/miniconda3/lib/python3.7/site-packages/dask/array/core.py in store(sources, targets, lock, regions, compute, return_stored, **kwargs)
    943 
    944         if compute:
--> 945             result.compute(**kwargs)
    946             return None
    947         else:

~/miniconda3/lib/python3.7/site-packages/dask/base.py in compute(self, **kwargs)
    164         dask.base.compute
    165         """
--> 166         (result,) = compute(self, traverse=False, **kwargs)
    167         return result
    168 

~/miniconda3/lib/python3.7/site-packages/dask/base.py in compute(*args, **kwargs)
    442         postcomputes.append(x.__dask_postcompute__())
    443 
--> 444     results = schedule(dsk, keys, **kwargs)
    445     return repack([f(r, *a) for r, (f, a) in zip(results, postcomputes)])
    446 

~/miniconda3/lib/python3.7/site-packages/distributed/client.py in get(self, dsk, keys, restrictions, loose_restrictions, resources, sync, asynchronous, direct, retries, priority, fifo_timeout, actors, **kwargs)
   2712             retries=retries,
   2713             user_priority=priority,
-> 2714             actors=actors,
   2715         )
   2716         packed = pack_data(keys, futures)

~/miniconda3/lib/python3.7/site-packages/distributed/client.py in _graph_to_futures(self, dsk, keys, restrictions, loose_restrictions, priority, user_priority, resources, retries, fifo_timeout, actors)
   2639                 {
   2640                     "op": "update-graph",
-> 2641                     "tasks": valmap(dumps_task, dsk),
   2642                     "dependencies": dependencies,
   2643                     "keys": list(map(tokey, keys)),

~/miniconda3/lib/python3.7/site-packages/cytoolz/dicttoolz.pyx in cytoolz.dicttoolz.valmap()

~/miniconda3/lib/python3.7/site-packages/cytoolz/dicttoolz.pyx in cytoolz.dicttoolz.valmap()

~/miniconda3/lib/python3.7/site-packages/distributed/worker.py in dumps_task(task)
   3356             return d
   3357         elif not any(map(_maybe_complex, task[1:])):
-> 3358             return {"function": dumps_function(task[0]), "args": warn_dumps(task[1:])}
   3359     return to_serialize(task)
   3360 

~/miniconda3/lib/python3.7/site-packages/distributed/worker.py in warn_dumps(obj, dumps, limit)
   3365 def warn_dumps(obj, dumps=pickle.dumps, limit=1e6):
   3366     """ Dump an object to bytes, warn if those bytes are large """
-> 3367     b = dumps(obj, protocol=4)
   3368     if not _warn_dumps_warned[0] and len(b) > limit:
   3369         _warn_dumps_warned[0] = True

~/miniconda3/lib/python3.7/site-packages/distributed/protocol/pickle.py in dumps(x, buffer_callback, protocol)
     58         try:
     59             buffers.clear()
---> 60             result = cloudpickle.dumps(x, **dump_kwargs)
     61         except Exception as e:
     62             logger.info("Failed to serialize %s. Exception: %s", x, e)

~/miniconda3/lib/python3.7/site-packages/cloudpickle/cloudpickle_fast.py in dumps(obj, protocol)
    100         with io.BytesIO() as file:
    101             cp = CloudPickler(file, protocol=protocol)
--> 102             cp.dump(obj)
    103             return file.getvalue()
    104 

~/miniconda3/lib/python3.7/site-packages/cloudpickle/cloudpickle_fast.py in dump(self, obj)
    561     def dump(self, obj):
    562         try:
--> 563             return Pickler.dump(self, obj)
    564         except RuntimeError as e:
    565             if "recursion" in e.args[0]:

~/miniconda3/lib/python3.7/pickle.py in dump(self, obj)
    435         if self.proto >= 4:
    436             self.framer.start_framing()
--> 437         self.save(obj)
    438         self.write(STOP)
    439         self.framer.end_framing()

~/miniconda3/lib/python3.7/pickle.py in save(self, obj, save_persistent_id)
    502         f = self.dispatch.get(t)
    503         if f is not None:
--> 504             f(self, obj) # Call unbound method with explicit self
    505             return
    506 

~/miniconda3/lib/python3.7/pickle.py in save_tuple(self, obj)
    787         write(MARK)
    788         for element in obj:
--> 789             save(element)
    790 
    791         if id(obj) in memo:

~/miniconda3/lib/python3.7/pickle.py in save(self, obj, save_persistent_id)
    522             reduce = getattr(obj, "__reduce_ex__", None)
    523             if reduce is not None:
--> 524                 rv = reduce(self.proto)
    525             else:
    526                 reduce = getattr(obj, "__reduce__", None)

TypeError: can't pickle _thread.lock objects

Steps to Reproduce

Following the ‘What’s New in 4.1’:

from dask.distributed import Client
client = Client()
import dask.array as da
array = da.random.random((1000, 1000))
from astropy.io import fits
hdu = fits.PrimaryHDU(data=array)
hdu.writeto('test_dask.fits', overwrite=True)

System Details

Darwin-19.6.0-x86_64-i386-64bit
Python 3.7.6 | packaged by conda-forge | (default, Jun  1 2020, 18:33:30) 
[Clang 9.0.1 ]
Numpy 1.19.2
astropy 4.2
Scipy 1.3.1
Matplotlib 3.3.1
Dask 2.17.2