**kwargs not working in mpi4py MPIPoolExecutor

Question

The mpi4py documentation claims that you can pass **kwargs through the MPIPoolExecutor, but I haven't been able to get it to work. I execute the following code:

import time
import socket
import numpy as np
from mpi4py.futures import MPIPoolExecutor
from mpi4py import MPI


def print_val(x, kwargsstr):
    time.sleep(1)
    msg = "x value: " + str(x) + \
        ", socket: " + str(socket.gethostname()) + \
        ", rank: " + str(MPI.COMM_WORLD.Get_rank()) + \
        ", time: " + str(np.round(time.time(), 1)) + \
        ", kwargs string: " + kwargsstr
    print(msg)

    return x


def main():

    kwarg = 'test'
    kwargs = {'kwargsstr':kwarg}

    with MPIPoolExecutor() as executor:
        x_v = []
        # for res in executor.map(print_val,
        #                         range(9), 9*[kwarg]):
        for res in executor.map(print_val,
                                range(9), **kwargs):
            x_v += [res]
        print(x_v)


if __name__ == '__main__':

    """
    run on command line with 1 scheduler and 2 workers:
    $ mpiexec -n 1 -usize 3 -machinefile hostfile.txt python mpi4py_kwargs.py
    """

    main()

via this command:

$ mpiexec -n 1 -usize 3 -machinefile hostfile.txt python mpi4py_kwargs.py

and get this error message:

TypeError: print_val() missing 1 required positional argument: 'kwargsstr'

Note that when the commented out portion in main is switched the code runs as expected.

Answer 1

I don't think map supports kwargs. Let's have a look at the source code

map indeed takes kwargs in its signature:

def map(self, fn, *iterables, **kwargs):

This is what the documentation says about the kwargs arguments:

Keyword Args:
  unordered: If ``True``, yield results out-of-order, as completed.

It is not obvious whether keyword arguments, excluding unordered, will be passed to the callables. We continue. The implementation is done using:

iterable = getattr(itertools, 'izip', zip)(*iterables)
return self.starmap(fn, iterable, **kwargs)

Thus, kwargs is forwarded to starmap:

def starmap(self, fn, iterable, timeout=None, chunksize=1, **kwargs):

with the same documentation

Keyword Args:
  unordered: If ``True``, yield results out-of-order, as completed.

The implementation looks as follows:

unordered = kwargs.pop('unordered', False)
if chunksize < 1:
    raise ValueError("chunksize must be >= 1.")
if chunksize == 1:
    return _starmap_helper(self.submit, fn, iterable,
                           timeout, unordered)
else:
    return _starmap_chunks(self.submit, fn, iterable,
                           timeout, unordered, chunksize)

Apparently, kwargs is not used any further, once unordered has been retrieved. Thus, map does not consider kwargs as callable kwargs.

As pointed out in my comment, for this particular example, you can provide the argument as non-keyword-argument, but positional argument, as well:

for res in executor.map(print_val, range(9), [kwarg] * 9):

as the second argument of map is documented as:

iterables: Iterables yielding positional arguments to be passed to
           the callable.

For people new to python:

>>> kwarg = 'test'
>>> [kwarg] * 9
['test', 'test', 'test', 'test', 'test', 'test', 'test', 'test', 'test']

where [kwarg] * 9 is a list implementing the iterables protocol.