`RuntimeError` when sending graph with `dask.array.ufunc.ufunc`s to scheduler and importing the wrapped `numpy` function
- Dominant language
- Python
- Stars
- 1.7k
- Forks
- 778
- Avg merge
- 2h 50m
- Merged PRs (30d)
- 3
Description
**Describe the issue**:
There's an odd pickle error that occurs when a `numpy` function that `dask.array` wraps in `dask.array.ufunc.ufunc`:
**Minimal Complete Verifiable Example**:
```python3
import dask
from numpy import exp
import dask.array as da
from distributed import Client, SubprocessCluster
if __name__ == "__main__":
client = Client(SubprocessCluster())
da.exp(da.from_array([1,2,3])).compute()
```
**Traceback**
```python3
Traceback (most recent call last):
File "/Users/hendrikmakait/projects/dask/distributed/distributed/scheduler.py", line 4659, in update_graph
graph = deserialize(graph_header, graph_frames).data
^^^^^^^^^^^
File "/Users/hendrikmakait/projects/dask/distributed/distributed/protocol/serialize.py", line 439, in deserialize
return loads(header, frames)
^^^^^^^^^^^^^^^^^
File "/Users/hendrikmakait/projects/dask/distributed/distributed/protocol/serialize.py", line 101, in pickle_loads
return pickle.loads(x, buffers=buffers)
^^^^^^^^^^^^^^^^^
File "/Users/hendrikmakait/projects/dask/distributed/distributed/protocol/pickle.py", line 94, in loads
return pickle.loads(x, buffers=buffers)
^^^^^^^^^^^^^^^
AttributeError: Can't get attribute 'exp' on
The above exception was the direct cause of the following exception:
Traceback (most recent call last):
File "/Users/hendrikmakait/projects/dask/distributed/sandbox.py", line 8, in
da.exp(da.from_array([1,2,3])).compute()
File "/opt/homebrew/Caskroom/mambaforge/base/envs/dask-distributed/lib/python3.12/site-packages/dask/base.py", line 358, in compute
(result,) = compute(self, traverse=False, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/opt/homebrew/Caskroom/mambaforge/base/envs/dask-distributed/lib/python3.12/site-packages/dask/base.py", line 644, in compute
results = schedule(dsk, keys, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/hendrikmakait/projects/dask/distributed/distributed/client.py", line 2243, in _gather
raise exception.with_traceback(traceback)
RuntimeError: Error during deserialization of the task graph. This frequently
occurs if the Scheduler and Client have different environments.
For more information, see
https://docs.dask.org/en/stable/deployment-considerations.html#consistent-software-environments
```
This seems to be reproducible with all `ufunc`s. For example, try replacing `exp` with `absolute` or `sin`. The `SubprocessCluster` is not needed, the scheduler just has to run in a different process (i.e., no process-local `LocalCluster`).
**Anything else we need to know?**:
This example starts working if we do either of the following:
* Remove `import dask`
* Remove `from numpy import exp`
* Replace `da.exp(...)` with a different ufunc that is not imported, e.g., `da.sin(...)`
**Environment**:
- Dask version: `main`
- Python version: 3.12
- Operating System: mac os 14.2.1
- Install method (conda, pip, source): source
Contributor guide
Assessment
This issue has not been assessed yet.