Python pickle AttributeError ‘module’ object has no attribute ‘foobar’

Ran into an interesting edge case with pickle this week. I had a producer task that was querying objects from a database, and pickling them plus a reference to a callback function to pass to worker tasks. Everything was working fine, but I was getting sick of logging into a Django shell to invoke the workers with test data. So I wrote a quick __main__ function in my task code to do the same thing.

import argparse
import pickle
from django.contrib.auth import user
from myproject.utils import my_task_func, worker


if __name__ == "__main__":
parser = argparse.ArgumentParser(description='Invoke the workers manually.')
parser.add_argument('user_id', type=int, help='User ID')
args = parser.parse_args()
user = User.objects.get(pk=args.user_id)
pickled_callback = pickle.dumps(my_task_func)
pickled_cb_args = pickle.dumps([user])
worker.delay(pickled_callback, pickled_cb_args)

To my surprise, this exact same code I had been typing into a shell was now throwing an exception when invoked from inside a __main__ function.

Traceback (most recent call last):
...
File "/usr/lib/python2.4/pickle.py", line 872, in load
dispatch[key](self)
File "/usr/lib/python2.4/pickle.py", line 1083, in load_inst
klass = self.find_class(module, name)
File "/usr/lib/python2.4/pickle.py", line 1140, in find_class
klass = getattr(mod, name)
AttributeError: 'module' object has no attribute 'my_task_func'

It turns out that when you pickle a function from inside a __main__ block, the module reference that it will be pickled with is __main__, not the actual module namespace. This is actually a known issue. In my case, it was easy to work-around; I simply put this code into a Django custom management command.

More discussion of this issue here.