Cray MPI Wrapper

wrapper.py

#!/usr/bin/env python
import os
import sys
import shlex
os.putenv('RANK', os.getenv('PMI_RANK'))
os.putenv('WORLD_SIZE', os.getenv('PMI_SIZE'))
os.putenv('LOCAL_RANK', os.getenv('PMI_LOCAL_RANK'))
os.putenv('LOCAL_SIZE', os.getenv('PMI_LOCAL_SIZE'))
argv = []
if os.getenv('PMI_RANK') == '0':
    argv = shlex.split(os.getenv('WRAPPER_CMD'))
argv += sys.argv[1:]
os.execvp(argv[0], argv)