How to pass an environmental variable with different values to spawned processes using MPI4PY.

mpi4py_self_spawn_multiple.py

#!/usr/bin/env python

"""
Self-launching MPI4PY script that uses spawning and sets an environmental 
variable to different values for each of the launched script instances..

Notes
-----
Requires an MPI implementation that supports dynamic process management.
"""

import os
import sys

from mpi4py import MPI

def worker():
    comm = MPI.Comm.Get_parent()
    size = comm.Get_size()
    rank = comm.Get_rank()
    name = MPI.Get_processor_name()

    print 'I am process %d of %d on %s [MYVAR=%s]' % \
        (rank, size, name, os.environ['MYVAR'])
    comm.Disconnect()

if __name__ == '__main__':
    script_file_name = os.path.basename(__file__)

    # The parent of the spawning process has a null communicator:
    if MPI.Comm.Get_parent() != MPI.COMM_NULL:
        worker()
    else:

        # Create lists of parameters:
        maxprocs = 3
        cmd_list = [sys.executable for i in xrange(maxprocs)]
        args_list = [script_file_name for i in xrange(maxprocs)]
        maxprocs_list = [1 for i in xrange(maxprocs)]
        info_list = [MPI.Info.Create() for i in xrange(maxprocs)]
        for i in xrange(maxprocs):
            info_list[i].Set('env', 'MYVAR=%s\n' % i)

        # Launch:
        comm = MPI.COMM_SELF.Spawn_multiple(cmd_list,
                                            args=args_list,
                                            maxprocs=maxprocs_list,
                                            info=info_list)
        comm.Disconnect()
        print 'I am the launcher'