"can't pickle thread.lock objects" when calling array.store with distributed

[rabernat]

I am trying to store a dask array using distributed. When I call store, I get an error "can't pickle thread.lock objects".

I originally was trying this in a much more complex context involving netCDF, xarray, etc. But I managed to come up with the following minimal example.

import numpy as np
import dask.array as da
from distributed import Client

def create_and_store_dask_array():
    shape = (10000, 1000)
    chunks = (1000, 1000)
    data = da.zeros(shape, chunks=chunks)
    store = np.memmap('test.memmap', mode='w+', dtype=data.dtype, shape=data.shape)
    data.store(store)
    print("Success!")

create_and_store_dask_array()
client = Client()
create_and_store_dask_array()

The first call works, but the second fails. The output is:

Success!
/home/rpa/.conda/envs/lagrangian_vorticity/lib/python2.7/site-packages/distributed/protocol/pickle.pyc - INFO - Failed to serialize (<function store at 0x7f0ee802f488>, (<functools.partial object at 0x7f0ec84f1418>, (1000, 1000)), (slice(2000, 3000, None), slice(0, 1000, None)), <thread.lock object at 0x7f0f2c715af0>)
Traceback (most recent call last):
  File "/home/rpa/.conda/envs/lagrangian_vorticity/lib/python2.7/site-packages/distributed/protocol/pickle.py", line 43, in dumps
    return cloudpickle.dumps(x, protocol=pickle.HIGHEST_PROTOCOL)
  File "/home/rpa/.conda/envs/lagrangian_vorticity/lib/python2.7/site-packages/cloudpickle/cloudpickle.py", line 706, in dumps
    cp.dump(obj)
  File "/home/rpa/.conda/envs/lagrangian_vorticity/lib/python2.7/site-packages/cloudpickle/cloudpickle.py", line 146, in dump
    return Pickler.dump(self, obj)
  File "/home/rpa/.conda/envs/lagrangian_vorticity/lib/python2.7/pickle.py", line 224, in dump
    self.save(obj)
  File "/home/rpa/.conda/envs/lagrangian_vorticity/lib/python2.7/pickle.py", line 286, in save
    f(self, obj) # Call unbound method with explicit self
  File "/home/rpa/.conda/envs/lagrangian_vorticity/lib/python2.7/pickle.py", line 568, in save_tuple
    save(element)
  File "/home/rpa/.conda/envs/lagrangian_vorticity/lib/python2.7/pickle.py", line 306, in save
    rv = reduce(self.proto)
TypeError: can't pickle thread.lock objects
/home/rpa/.conda/envs/lagrangian_vorticity/lib/python2.7/site-packages/distributed/protocol/core.pyc - CRITICAL - Failed to Serialize
Traceback (most recent call last):
  File "/home/rpa/.conda/envs/lagrangian_vorticity/lib/python2.7/site-packages/distributed/protocol/core.py", line 43, in dumps
    for key, value in data.items()
  File "/home/rpa/.conda/envs/lagrangian_vorticity/lib/python2.7/site-packages/distributed/protocol/core.py", line 44, in <dictcomp>
    if type(value) is Serialize}
  File "/home/rpa/.conda/envs/lagrangian_vorticity/lib/python2.7/site-packages/distributed/protocol/serialize.py", line 106, in serialize
    header, frames = {}, [pickle.dumps(x)]
  File "/home/rpa/.conda/envs/lagrangian_vorticity/lib/python2.7/site-packages/distributed/protocol/pickle.py", line 43, in dumps
    return cloudpickle.dumps(x, protocol=pickle.HIGHEST_PROTOCOL)
  File "/home/rpa/.conda/envs/lagrangian_vorticity/lib/python2.7/site-packages/cloudpickle/cloudpickle.py", line 706, in dumps
    cp.dump(obj)
  File "/home/rpa/.conda/envs/lagrangian_vorticity/lib/python2.7/site-packages/cloudpickle/cloudpickle.py", line 146, in dump
    return Pickler.dump(self, obj)
  File "/home/rpa/.conda/envs/lagrangian_vorticity/lib/python2.7/pickle.py", line 224, in dump
    self.save(obj)
  File "/home/rpa/.conda/envs/lagrangian_vorticity/lib/python2.7/pickle.py", line 286, in save
    f(self, obj) # Call unbound method with explicit self
  File "/home/rpa/.conda/envs/lagrangian_vorticity/lib/python2.7/pickle.py", line 568, in save_tuple
    save(element)
  File "/home/rpa/.conda/envs/lagrangian_vorticity/lib/python2.7/pickle.py", line 306, in save
    rv = reduce(self.proto)
TypeError: can't pickle thread.lock objects

Versions:

import dask
print dask.__version__
import distributed
print distributed.__version__
>>> 0.12.0
>>> 1.14.3

[rabernat]

Seems related to dask/dask#1683.

[mrocklin]

Can you try the following

# data.store(store)

from dask.utils import SerializableLock
lock = SerializableLock()
data.store(store, lock=lock)

[rabernat]

It works with SerializableLock. Thanks @mrocklin!

Is there a way to make this the default?

[rabernat]

p.s. I had to update to dask. 0.13.0 in order to be able to import the SerializableLock.

[mrocklin]

We could make it the default, yes. I suppose it doesn't have much overhead over standard locks.

[mrocklin]

This just came in. You could ping there and see if he's interested in implementing this as well: dask/knit#60

[rabernat]

So with SerializableLock it works at the dask level. But now I would like to wrap this array in an xarray Dataset and call to_netcdf. Unfortunately that function does not allow me to pass a custom lock. It appears that calls to dask.array.store automatically get passed a threading.Lock() object:
https://github.com/pydata/xarray/blob/master/xarray/backends/common.py#L153

Should I raise this issue over at xarray then?

(I must admit I don't see the relevance of dask/knit#60 to this issue.)

[mrocklin]

Sorry, wrong copy-paste: dask/dask#1879

[mrocklin]

Or you could also implement this yourself. Should be an easy change and it'd be good to have your name in the authors list :) (or I can do it, but that's boring)

[mrocklin]

We built SerializableLock with XArray in mind. I think that they started using it as well in some cases.

cc @shoyer

[rabernat]

Ah, it looks like there is already an xarray PR in progress that would address this: pydata/xarray#1179

[pelson]

Just ran into this too. Is it feasible to extend cloudpickle to dispatch thread.lock, rather than retro-fitting SerializableLock usage?

[mrocklin]

I think that cloudpickle should probably still err when trying to serialize a lock. This is usually an unsafe thing to do I think.