Failing to deserialize user function

[rjzamora]

Ran into this error in RAPIDS CI, but no GPU is needed to reproduce:

import dask.dataframe as dd
from dask.distributed import LocalCluster, Client, wait

if __name__ == "__main__":
    with LocalCluster() as cluster:
        client = Client(cluster)
        ddf = dd.from_dict({"a": range(10)}, 1)
        def func(df):
            return (df + 5) * 3 - 11
        wait(ddf.map_partitions(func).persist())

...
  File "/datasets/rzamora/miniconda3/envs/cudf_2402_dev/lib/python3.10/site-packages/distributed/protocol/core.py", line 172, in _decode_default
    return pickle.loads(sub_header["pickled-obj"], buffers=sub_frames)
  File "/datasets/rzamora/miniconda3/envs/cudf_2402_dev/lib/python3.10/site-packages/distributed/protocol/pickle.py", line 99, in loads
    return pickle.loads(x, buffers=buffers)
AttributeError: Can't get attribute 'func' on <module '__mp_main__' from '/home/nfs/rzamora/workspace/cudf-24.02/bug.py'>

I suspect that this was cause by #8447, but haven't confirmed this yet.

[hendrikmakait]

It looks like #8443 is to blame.

[rjzamora]

It looks like #8443 is to blame.

Cool, was just about to say #8447 is not the problem.

[hendrikmakait]

Feel free to dig deeper into this, otherwise I'll have a go at it tomorrow to get this fixed before the release.