urgent: uproot.dask(..., open_files=True) not usable with dask distributed client
lgray opened this issue · comments
uproot version: 5.2.0rc4 (but also exists in 5.2.0rc1-3)
@lobis @iasonkrom
The following code:
import dask
from distributed import Client
import uproot
if __name__ == "__main__":
with Client() as _:
events = uproot.dask("tests/samples/nano_dy.root:Events")
print(dask.compute(events.Muon_pt))fails with a very lengthy error message, the gist of which is:
Traceback (most recent call last):
File "/Users/lgray/miniforge3/envs/coffea-rc-test/lib/python3.10/site-packages/distributed/protocol/pickle.py", line 63, in dumps
result = pickle.dumps(x, **dump_kwargs)
AttributeError: Can't pickle local object 'unpack_collections.<locals>.repack'
During handling of the above exception, another exception occurred:
Traceback (most recent call last):
File "/Users/lgray/miniforge3/envs/coffea-rc-test/lib/python3.10/site-packages/distributed/protocol/pickle.py", line 68, in dumps
pickler.dump(x)
AttributeError: Can't pickle local object 'unpack_collections.<locals>.repack'
During handling of the above exception, another exception occurred:
Traceback (most recent call last):
File "/Users/lgray/miniforge3/envs/coffea-rc-test/lib/python3.10/site-packages/distributed/protocol/pickle.py", line 81, in dumps
result = cloudpickle.dumps(x, **dump_kwargs)
File "/Users/lgray/miniforge3/envs/coffea-rc-test/lib/python3.10/site-packages/cloudpickle/cloudpickle.py", line 1479, in dumps
cp.dump(obj)
File "/Users/lgray/miniforge3/envs/coffea-rc-test/lib/python3.10/site-packages/cloudpickle/cloudpickle.py", line 1245, in dump
return super().dump(obj)
File "/Users/lgray/miniforge3/envs/coffea-rc-test/lib/python3.10/site-packages/uproot/source/fsspec.py", line 68, in __getstate__
state.pop("_fh")
KeyError: '_fh'
2023-12-13 08:39:47,143 - distributed.protocol.core - CRITICAL - Failed to Serialize
Traceback (most recent call last):
File "/Users/lgray/miniforge3/envs/coffea-rc-test/lib/python3.10/site-packages/distributed/protocol/pickle.py", line 63, in dumps
result = pickle.dumps(x, **dump_kwargs)
AttributeError: Can't pickle local object 'unpack_collections.<locals>.repack'
The key error above appears to be causing the pickling error that gets reported first. A bit confusing but makes sense in the end.
If we use open_files=False:
import dask
from distributed import Client
import uproot
if __name__ == "__main__":
with Client() as _:
events = uproot.dask("tests/samples/nano_dy.root:Events", open_files=False)
print(dask.compute(events.Muon_pt))
it successfully executes:
(coffea-rc-test) lgray@Lindseys-MacBook-Pro coffea % python uproot_distributed_test.py
(<Array [[], [], [], [], [...], ..., [], [], [], []] type='40 * var * float32'>,)
This should be fixed before we release uproot 5.2.0.
We should also add some tests with distributed to catch stuff like this in the future!
This will be closed by #1022 once it's merged into main
This was not resolved:
Traceback (most recent call last):
File /lib/python3.11/site-packages/distributed/protocol/pickle.py", line 63, in dumps
result = pickle.dumps(x, **dump_kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
TypeError: cannot pickle '_thread.lock' object
During handling of the above exception, another exception occurred:
Traceback (most recent call last):
File /lib/python3.11/site-packages/distributed/protocol/pickle.py", line 68, in dumps
pickler.dump(x)
TypeError: cannot pickle '_thread.lock' object
During handling of the above exception, another exception occurred:
Traceback (most recent call last):
File /lib/python3.11/site-packages/distributed/protocol/serialize.py", line 353, in serialize
header, frames = dumps(x, context=context) if wants_context else dumps(x)
^^^^^^^^^^^^^^^^^^^^^^^^^
File /lib/python3.11/site-packages/distributed/protocol/serialize.py", line 76, in pickle_dumps
frames[0] = pickle.dumps(
^^^^^^^^^^^^^
File /lib/python3.11/site-packages/distributed/protocol/pickle.py", line 81, in dumps
result = cloudpickle.dumps(x, **dump_kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File /lib/python3.11/site-packages/cloudpickle/cloudpickle_fast.py", line 73, in dumps
cp.dump(obj)
File /lib/python3.11/site-packages/cloudpickle/cloudpickle_fast.py", line 632, in dump
return Pickler.dump(self, obj)
^^^^^^^^^^^^^^^^^^^^^^^
TypeError: cannot pickle '_thread.lock' object
The above exception was the direct cause of the following exception:
Traceback (most recent call last):
File "test.py", line 11, in <module>
print(dask.compute(events["branch"]))
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File /lib/python3.11/site-packages/dask/base.py", line 628, in compute
results = schedule(dsk, keys, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File /lib/python3.11/site-packages/distributed/protocol/serialize.py", line 379, in serialize
raise TypeError(msg, str_x) from exc
TypeError: ('Could not serialize object of type HighLevelGraph', '<ToPickle: HighLevelGraph with 1 layers.\n<dask.highlevelgraph.HighLevelGraph object at 0x7f19ff0041d0>\n 0. branch-f4983d883434c3c122dafd5d2e31979f\n>')
This seems like a user error, _thread.lock should never show up as an object in a task graph to begin with.
Please provide a reproducer!
As well as versions of awkward/dask-awkward/uproot.
@lgray please checkout https://github.com/veprbl/uproot-issue1063
This is reproducible with
awkward.__version__ 2.5.1
dask.__version__ 2023.12.0
dask_awkward.__version__ 2023.12.2
distributed.__version__ 2023.12.0
uproot.__version__ 5.2.0
Cool, glad it's not user error. We'll see what's up with it.
This is now tracked in #1194