modules using multiprocessing fail with mitogen: PicklingError: Can't pickle <function ...>

[asheplyakov]

Here is a simple module (called countmp)

#!/usr/bin/env python3

import multiprocessing as mp

from ansible.module_utils.basic import AnsibleModule


def fmt(x):
    return str(x)


def count(start, end):
    with mp.Pool() as pool:
        return list(pool.imap(fmt, range(start, end + 1)))


def main():
    module = AnsibleModule(
        argument_spec=dict(
            start=dict(type='int', required=False, default=0),
            end=dict(type='int', required=True),
        )
    )
    start = module.params['start']
    end = module.params['end']
    result = {
        'start': start,
        'end': end,
        'result': count(start, end),
        'changed': False,
    }
    module.exit_json(**result)


if __name__ == '__main__':
    main()

It works just fine without mitogen:

$ ansible -m countmp -a 'start=1 end=5' localhost
localhost | SUCCESS => {
    "changed": false,
    "end": 5,
    "result": [
        "1",
        "2",
        "3",
        "4",
        "5"
    ],
    "start": 1
}

However it fails with mitogen:

$ ansible -m countmp -a 'start=1 end=5' localhost
An exception occurred during task execution. To see the full traceback, use -vvv. The error was: _pickle.PicklingError: Can't pickle <function fmt at 0x7f7c62335940>: attribute lookup fmt on __main__ failed
localhost | FAILED! => {
    "changed": false,
    "module_stderr": "Traceback (most recent call last):\n  File \"master:/home/asheplyakov/hacking/mitogen/mitogen/ansible_mitogen/runner.py\", line 975, in _run\n    self._run_code(code, mod)\n  File \"master:/home/asheplyakov/hacking/mitogen/mitogen/ansible_mitogen/runner.py\", line 939, in _run_code\n    exec(code, vars(mod))\n  File \"master:/home/asheplyakov/hacking/mitogen/library/countmp.py\", line 36, in <module>\n  File \"master:/home/asheplyakov/hacking/mitogen/library/countmp.py\", line 29, in main\n  File \"master:/home/asheplyakov/hacking/mitogen/library/countmp.py\", line 14, in count\n  File \"/usr/lib64/python3.9/multiprocessing/pool.py\", line 870, in next\n    raise value\n  File \"/usr/lib64/python3.9/multiprocessing/pool.py\", line 537, in _handle_tasks\n    put(task)\n  File \"/usr/lib64/python3.9/multiprocessing/connection.py\", line 211, in send\n    self._send_bytes(_ForkingPickler.dumps(obj))\n  File \"/usr/lib64/python3.9/multiprocessing/reduction.py\", line 51, in dumps\n    cls(buf, protocol).dump(obj)\n_pickle.PicklingError: Can't pickle <function fmt at 0x7f7c62335940>: attribute lookup fmt on __main__ failed\n",
    "module_stdout": "",
    "msg": "MODULE FAILURE\nSee stdout/stderr for the exact error",
    "rc": 1
}

The problem is that multiprocessing.Pool needs to pickle arguments to send them to worker processes. However pickling works only for functions defined in the __main__ module. Mitogen subtly changes the context of execution, so the fmt function is not defined within the __main__ module any more, hence the error.

$ ansible --version
ansible 2.9.27
  config file = /home/asheplyakov/hacking/mitogen/ansible.cfg
  configured module search path = ['/home/asheplyakov/hacking/mitogen/library']
  ansible python module location = /usr/lib/python3/site-packages/ansible
  executable location = /usr/bin/ansible
  python version = 3.9.6 (default, Jun 29 2021, 10:42:27) [GCC 10.2.1 20210313 (ALT Sisyphus 10.2.1-alt3)]

$ cat ansible.cfg
[defaults]
strategy_plugins = mitogen/ansible_mitogen/plugins/strategy
strategy = mitogen_linear
library = library