How to use twiggy for logging in both parent and MPI-spawned processes.

mpi4py_spawn_twiggy.py

#!/usr/bin/bin python

"""
How to use twiggy for logging in both parent and MPI-spawned processes.
"""

import sys

# Need to use dill to serialize the twiggy emitters because pickle can't handle
# them:
import dill
from mpi4py import MPI
MPI.pickle.dumps = dill.dumps
MPI.pickle.loads = dill.loads

import twiggy

rank = MPI.COMM_WORLD.Get_rank()

intercomm = MPI.Comm.Get_parent()
if intercomm == MPI.COMM_NULL:
    output = twiggy.outputs.StreamOutput(twiggy.formats.line_format,
                                         stream=sys.stdout)
    emitters = twiggy.filters.Emitter(twiggy.levels.DEBUG,
                                      True, output)
    twiggy.emitters['*'] = emitters
    logger = twiggy.log.name('parent')
    logger.info('in parent')
    maxprocs = 2
    comm = MPI.COMM_SELF.Spawn(sys.executable,
                               args=[__file__],
                               maxprocs=maxprocs)
    for i in xrange(maxprocs):
        comm.send(twiggy.emitters, i)
else:
    emitters = intercomm.recv()
    # The transmitted emitters must be explicitly associated with their keys in
    # the child twiggy.emitters dict:
    twiggy.emitters.update(emitters)
    logger = twiggy.log.name('child %s' % rank)
    logger.info('in child')