Distributed: Unable to serialize numba function if directly imported

[aktech]

What happened:
It seems distributed is not able to numba_function when it is imported directly. I understand this is a nested function, but it works when accessed via another_module.numba_function

What you expected to happen:
Distributed being able to serialize in both cases

Minimal Complete Verifiable Example:

Reproducing it involves two files:

• File: another_module.py

from numba import guvectorize

@guvectorize(["(float64, float64[:])"], '()->()')
def numba_function(x, res):
    res[0] = x

• File: distributed_bug.py

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

import another_module
from another_module import numba_function


def foo_works(x):
    def bar(f):
        return another_module.numba_function(f)
    out = x.map_blocks(bar)
    return out


def foo_does_not_works(x):
    def bar(f):
        return numba_function(f)
    out = x.map_blocks(bar)
    return out


if __name__ == '__main__':
    x = da.arange(4)
    client = Client()
    print("foo works:")
    print(foo_works(x).compute())
    print("foo doesn't works:")
    foo_does_not_works(x).compute()

Run the above file:

 python distributed_bug.py

Anything else we need to know?:
I doubt that it use to work in an earlier version, although I don't recall the exact version.

Environment:

# packages in environment at /Users/aktech/anaconda3/envs/distributed-bug:
dask                      2021.3.0           pyhd8ed1ab_0    conda-forge
dask-core                 2021.3.0           pyhd8ed1ab_0    conda-forge
distributed               2021.3.0         py38h50d1736_0    conda-forge
numba                     0.53.0           py38hb2f4e1b_0

• Dask version: 2021.3.0
• Python version: 3.8.8
• Operating System: Mac OS 11.2.3 (20D91)
• Install method (conda, pip, source): conda

Exception

Exception

``` foo works: [0. 1. 2. 3.] foo doesn't works: distributed.protocol.core - CRITICAL - Failed to deserialize Traceback (most recent call last): File "/Users/aktech/anaconda3/envs/distributed-bug/lib/python3.8/site-packages/distributed/protocol/core.py", line 130, in loads value = merge_and_deserialize(head, fs, deserializers=deserializers) File "/Users/aktech/anaconda3/envs/distributed-bug/lib/python3.8/site-packages/distributed/protocol/serialize.py", line 450, in merge_and_deserialize return deserialize(header, merged_frames, deserializers=deserializers) File "/Users/aktech/anaconda3/envs/distributed-bug/lib/python3.8/site-packages/distributed/protocol/serialize.py", line 384, in deserialize return loads(header, frames) File "/Users/aktech/anaconda3/envs/distributed-bug/lib/python3.8/site-packages/distributed/protocol/serialize.py", line 164, in serialization_error_loads raise TypeError(msg) TypeError: Could not serialize object of type tuple. Traceback (most recent call last): File "/Users/aktech/anaconda3/envs/distributed-bug/lib/python3.8/site-packages/distributed/protocol/pickle.py", line 49, in dumps result = pickle.dumps(x, **dump_kwargs) AttributeError: Can't pickle local object 'foo_does_not_works..bar' During handling of the above exception, another exception occurred: Traceback (most recent call last): File "/Users/aktech/anaconda3/envs/distributed-bug/lib/python3.8/site-packages/distributed/protocol/serialize.py", line 307, in serialize header, frames = dumps(x, context=context) if wants_context else dumps(x) File "/Users/aktech/anaconda3/envs/distributed-bug/lib/python3.8/site-packages/distributed/protocol/serialize.py", line 58, in pickle_dumps frames[0] = pickle.dumps( File "/Users/aktech/anaconda3/envs/distributed-bug/lib/python3.8/site-packages/distributed/protocol/pickle.py", line 60, in dumps result = cloudpickle.dumps(x, **dump_kwargs) File "/Users/aktech/anaconda3/envs/distributed-bug/lib/python3.8/site-packages/cloudpickle/cloudpickle_fast.py", line 73, in dumps cp.dump(obj) File "/Users/aktech/anaconda3/envs/distributed-bug/lib/python3.8/site-packages/cloudpickle/cloudpickle_fast.py", line 563, in dump return Pickler.dump(self, obj) _pickle.PicklingError: Can't pickle : it's not the same object as __main__.numba_function distributed.core - ERROR - Could not serialize object of type tuple. Traceback (most recent call last): File "/Users/aktech/anaconda3/envs/distributed-bug/lib/python3.8/site-packages/distributed/protocol/pickle.py", line 49, in dumps result = pickle.dumps(x, **dump_kwargs) AttributeError: Can't pickle local object 'foo_does_not_works..bar' During handling of the above exception, another exception occurred: Traceback (most recent call last): File "/Users/aktech/anaconda3/envs/distributed-bug/lib/python3.8/site-packages/distributed/protocol/serialize.py", line 307, in serialize header, frames = dumps(x, context=context) if wants_context else dumps(x) File "/Users/aktech/anaconda3/envs/distributed-bug/lib/python3.8/site-packages/distributed/protocol/serialize.py", line 58, in pickle_dumps frames[0] = pickle.dumps( File "/Users/aktech/anaconda3/envs/distributed-bug/lib/python3.8/site-packages/distributed/protocol/pickle.py", line 60, in dumps result = cloudpickle.dumps(x, **dump_kwargs) File "/Users/aktech/anaconda3/envs/distributed-bug/lib/python3.8/site-packages/cloudpickle/cloudpickle_fast.py", line 73, in dumps cp.dump(obj) File "/Users/aktech/anaconda3/envs/distributed-bug/lib/python3.8/site-packages/cloudpickle/cloudpickle_fast.py", line 563, in dump return Pickler.dump(self, obj) _pickle.PicklingError: Can't pickle : it's not the same object as __main__.numba_function Traceback (most recent call last): File "/Users/aktech/anaconda3/envs/distributed-bug/lib/python3.8/site-packages/distributed/core.py", line 554, in handle_stream msgs = await comm.read() File "/Users/aktech/anaconda3/envs/distributed-bug/lib/python3.8/site-packages/distributed/comm/tcp.py", line 217, in read msg = await from_frames( File "/Users/aktech/anaconda3/envs/distributed-bug/lib/python3.8/site-packages/distributed/comm/utils.py", line 80, in from_frames res = _from_frames() File "/Users/aktech/anaconda3/envs/distributed-bug/lib/python3.8/site-packages/distributed/comm/utils.py", line 63, in _from_frames return protocol.loads( File "/Users/aktech/anaconda3/envs/distributed-bug/lib/python3.8/site-packages/distributed/protocol/core.py", line 130, in loads value = merge_and_deserialize(head, fs, deserializers=deserializers) File "/Users/aktech/anaconda3/envs/distributed-bug/lib/python3.8/site-packages/distributed/protocol/serialize.py", line 450, in merge_and_deserialize return deserialize(header, merged_frames, deserializers=deserializers) File "/Users/aktech/anaconda3/envs/distributed-bug/lib/python3.8/site-packages/distributed/protocol/serialize.py", line 384, in deserialize return loads(header, frames) File "/Users/aktech/anaconda3/envs/distributed-bug/lib/python3.8/site-packages/distributed/protocol/serialize.py", line 164, in serialization_error_loads raise TypeError(msg) TypeError: Could not serialize object of type tuple. Traceback (most recent call last): File "/Users/aktech/anaconda3/envs/distributed-bug/lib/python3.8/site-packages/distributed/protocol/pickle.py", line 49, in dumps result = pickle.dumps(x, **dump_kwargs) AttributeError: Can't pickle local object 'foo_does_not_works..bar' During handling of the above exception, another exception occurred: Traceback (most recent call last): File "/Users/aktech/anaconda3/envs/distributed-bug/lib/python3.8/site-packages/distributed/protocol/serialize.py", line 307, in serialize header, frames = dumps(x, context=context) if wants_context else dumps(x) File "/Users/aktech/anaconda3/envs/distributed-bug/lib/python3.8/site-packages/distributed/protocol/serialize.py", line 58, in pickle_dumps frames[0] = pickle.dumps( File "/Users/aktech/anaconda3/envs/distributed-bug/lib/python3.8/site-packages/distributed/protocol/pickle.py", line 60, in dumps result = cloudpickle.dumps(x, **dump_kwargs) File "/Users/aktech/anaconda3/envs/distributed-bug/lib/python3.8/site-packages/cloudpickle/cloudpickle_fast.py", line 73, in dumps cp.dump(obj) File "/Users/aktech/anaconda3/envs/distributed-bug/lib/python3.8/site-packages/cloudpickle/cloudpickle_fast.py", line 563, in dump return Pickler.dump(self, obj) _pickle.PicklingError: Can't pickle : it's not the same object as __main__.numba_function distributed.worker - ERROR - Could not serialize object of type tuple. Traceback (most recent call last): File "/Users/aktech/anaconda3/envs/distributed-bug/lib/python3.8/site-packages/distributed/protocol/pickle.py", line 49, in dumps result = pickle.dumps(x, **dump_kwargs) AttributeError: Can't pickle local object 'foo_does_not_works..bar' During handling of the above exception, another exception occurred: Traceback (most recent call last): File "/Users/aktech/anaconda3/envs/distributed-bug/lib/python3.8/site-packages/distributed/protocol/serialize.py", line 307, in serialize header, frames = dumps(x, context=context) if wants_context else dumps(x) File "/Users/aktech/anaconda3/envs/distributed-bug/lib/python3.8/site-packages/distributed/protocol/serialize.py", line 58, in pickle_dumps frames[0] = pickle.dumps( File "/Users/aktech/anaconda3/envs/distributed-bug/lib/python3.8/site-packages/distributed/protocol/pickle.py", line 60, in dumps result = cloudpickle.dumps(x, **dump_kwargs) File "/Users/aktech/anaconda3/envs/distributed-bug/lib/python3.8/site-packages/cloudpickle/cloudpickle_fast.py", line 73, in dumps cp.dump(obj) File "/Users/aktech/anaconda3/envs/distributed-bug/lib/python3.8/site-packages/cloudpickle/cloudpickle_fast.py", line 563, in dump return Pickler.dump(self, obj) _pickle.PicklingError: Can't pickle : it's not the same object as __main__.numba_function Traceback (most recent call last): File "/Users/aktech/anaconda3/envs/distributed-bug/lib/python3.8/site-packages/distributed/worker.py", line 990, in handle_scheduler await self.handle_stream( File "/Users/aktech/anaconda3/envs/distributed-bug/lib/python3.8/site-packages/distributed/core.py", line 554, in handle_stream msgs = await comm.read() File "/Users/aktech/anaconda3/envs/distributed-bug/lib/python3.8/site-packages/distributed/comm/tcp.py", line 217, in read msg = await from_frames( File "/Users/aktech/anaconda3/envs/distributed-bug/lib/python3.8/site-packages/distributed/comm/utils.py", line 80, in from_frames res = _from_frames() File "/Users/aktech/anaconda3/envs/distributed-bug/lib/python3.8/site-packages/distributed/comm/utils.py", line 63, in _from_frames return protocol.loads( File "/Users/aktech/anaconda3/envs/distributed-bug/lib/python3.8/site-packages/distributed/protocol/core.py", line 130, in loads value = merge_and_deserialize(head, fs, deserializers=deserializers) File "/Users/aktech/anaconda3/envs/distributed-bug/lib/python3.8/site-packages/distributed/protocol/serialize.py", line 450, in merge_and_deserialize return deserialize(header, merged_frames, deserializers=deserializers) File "/Users/aktech/anaconda3/envs/distributed-bug/lib/python3.8/site-packages/distributed/protocol/serialize.py", line 384, in deserialize return loads(header, frames) File "/Users/aktech/anaconda3/envs/distributed-bug/lib/python3.8/site-packages/distributed/protocol/serialize.py", line 164, in serialization_error_loads raise TypeError(msg) TypeError: Could not serialize object of type tuple. Traceback (most recent call last): File "/Users/aktech/anaconda3/envs/distributed-bug/lib/python3.8/site-packages/distributed/protocol/pickle.py", line 49, in dumps result = pickle.dumps(x, **dump_kwargs) AttributeError: Can't pickle local object 'foo_does_not_works..bar' During handling of the above exception, another exception occurred: Traceback (most recent call last): File "/Users/aktech/anaconda3/envs/distributed-bug/lib/python3.8/site-packages/distributed/protocol/serialize.py", line 307, in serialize header, frames = dumps(x, context=context) if wants_context else dumps(x) File "/Users/aktech/anaconda3/envs/distributed-bug/lib/python3.8/site-packages/distributed/protocol/serialize.py", line 58, in pickle_dumps frames[0] = pickle.dumps( File "/Users/aktech/anaconda3/envs/distributed-bug/lib/python3.8/site-packages/distributed/protocol/pickle.py", line 60, in dumps result = cloudpickle.dumps(x, **dump_kwargs) File "/Users/aktech/anaconda3/envs/distributed-bug/lib/python3.8/site-packages/cloudpickle/cloudpickle_fast.py", line 73, in dumps cp.dump(obj) File "/Users/aktech/anaconda3/envs/distributed-bug/lib/python3.8/site-packages/cloudpickle/cloudpickle_fast.py", line 563, in dump return Pickler.dump(self, obj) _pickle.PicklingError: Can't pickle : it's not the same object as __main__.numba_function tornado.application - ERROR - Exception in callback functools.partial(>, exception=TypeError('Could not serialize object of type tuple.\nTraceback (most recent call last):\n File "/Users/aktech/anaconda3/envs/distributed-bug/lib/python3.8/site-packages/distributed/protocol/pickle.py", line 49, in dumps\n result = pickle.dumps(x, **dump_kwargs)\nAttributeError: Can\'t pickle local object \'foo_does_not_works..bar\'\n\nDuring handling of the above exception, another exception occurred:\n\nTraceback (most recent call last):\n File "/Users/aktech/anaconda3/envs/distributed-bug/lib/python3.8/site-packages/distributed/protocol/serialize.py", line 307, in serialize\n header, frames = dumps(x, context=context) if wants_context else dumps(x)\n File "/Users/aktech/anaconda3/envs/distributed-bug/lib/python3.8/site-packages/distributed/protocol/serialize.py", line 58, in pickle_dumps\n frames[0] = pickle.dumps(\n File "/Users/aktech/anaconda3/envs/distributed-bug/lib/python3.8/site-packages/distributed/protocol/pickle.py", line 60, in dumps\n result = cloudpickle.dumps(x, **dump_kwargs)\n File "/Users/aktech/anaconda3/envs/distributed-bug/lib/python3.8/site-packages/cloudpickle/cloudpickle_fast.py", line 73, in dumps\n cp.dump(obj)\n File "/Users/aktech/anaconda3/envs/distributed-bug/lib/python3.8/site-packages/cloudpickle/cloudpickle_fast.py" , line 563, in dump\n return Pickler.dump(self, obj)\n_pickle.PicklingError: Can\'t pickle : it\'s not the same object as __main__.numba_function\n')>) \```

[jakirkham]

What happens if both are round-tripped through cloudpickle?

[jrbourbeau]

@aktech I wonder if you're running into this issue https://github.com/numba/numba/issues/6821. Trying numba < 0.53 might help

[aktech]

What happens if both are round-tripped through cloudpickle?

@jakirkham I haven't tried that yet

@aktech I wonder if you're running into this issue numba/numba#6821. Trying numba < 0.53 might help

Ah, looks like. I'll confirm after running on a downgraded Numba.

[jrbourbeau]

Just checking in here, @aktech did rolling back numba help? Is it safe to close this issue?

[aktech]

I haven't tried it yet, I shall report back (and close it) by the end of today hopefully.

[jrbourbeau]

Sounds good -- no rush, I just wanted to see how things were going : )

[aktech]

Thanks everyone it works on numba==0.52