import contextlib
import logging
import os
from datetime import datetime

from teuthology import misc as teuthology

log = logging.getLogger(__name__)

def task(ctx, config):
    """
    Setup MPI and execute commands

    Example that starts an MPI process on specific clients::

        tasks:
        - ceph:
        - ceph-fuse: [client.0, client.1]
        - mpi: 
            nodes: [client.0, client.1]
            exec: ior ...

    Example that starts MPI processes on all clients::

        tasks:
        - ceph:
        - ceph-fuse:
        - mpi:
            exec: ior ...

    Example that starts MPI processes on all roles::

        tasks:
        - ceph:
        - mpi:
            nodes: all
            exec: ...

    """
    assert isinstance(config, dict), 'task mpi got invalid config'
    assert 'exec' in config, 'task mpi got invalid config, missing exec'
    mpiexec = config['exec']
    hosts = []
    remotes = []
    master_remote = None
    if 'nodes' in config:
        if isinstance(config['nodes'], basestring) and config['nodes'] == 'all':
            for role in  teuthology.all_roles(ctx.cluster):
                (remote,) = ctx.cluster.only(role).remotes.iterkeys()
                ip,port = remote.ssh.get_transport().getpeername()
                hosts.append(ip)
                remotes.append(remote)
            (master_remote,) = ctx.cluster.only(config['nodes'][0]).remotes.iterkeys()
        elif isinstance(config['nodes'], list):
            for role in config['nodes']:
                (remote,) = ctx.cluster.only(role).remotes.iterkeys()
                ip,port = remote.ssh.get_transport().getpeername()
                hosts.append(ip)
                remotes.append(remote)
            (master_remote,) = ctx.cluster.only(config['nodes'][0]).remotes.iterkeys()
    else:
        roles = ['client.{id}'.format(id=id_) for id_ in teuthology.all_roles_of_type(ctx.cluster, 'client')]
        (master_remote,) = ctx.cluster.only(roles[0]).remotes.iterkeys()
        for role in roles:
            (remote,) = ctx.cluster.only(role).remotes.iterkeys()
            ip,port = remote.ssh.get_transport().getpeername()
            hosts.append(ip)
            remotes.append(remote)

    log.info('mpi rank 0 is: {name}'.format(name=master_remote.name))

    # write out the mpi hosts file
    log.info('mpi nodes: [%s]' % (', '.join(hosts)))
    hostfiledata = '\n'.join(hosts) + '\n'
    teuthology.write_file(remote=master_remote, path='/tmp/cephtest/mpi-hosts', data='\n'.join(hosts))
    log.info('mpiexec on {name}: {cmd}'.format(name=master_remote.name, cmd=mpiexec))
    args=['mpiexec', '-f', '/tmp/cephtest/mpi-hosts']
    args.extend(mpiexec.split(' '))
    master_remote.run(args=args, )
    log.info('mpi task completed')
    master_remote.run(args=['rm', '/tmp/cephtest/mpi-hosts'])

@@ -339,32 +339,52 @@
 nt.0, client.1%5D%0A
+        - ssh_keys:%0A
         - mpi: %0A
@@ -383,16 +383,16 @@
 - mpi: %0A
-
         
@@ -554,24 +554,44 @@
  ceph-fuse:%0A
+        - ssh_keys:%0A
         - mp
@@ -697,32 +697,52 @@
         - ceph:%0A
+        - ssh_keys:%0A
         - mpi:%0A 
@@ -763,16 +763,16 @@
 es: all%0A
-
         
@@ -782,24 +782,423 @@
  exec: ...%0A%0A
+    Example that specifies a working directory for MPI processes:%0A%0A        tasks:%0A        - ceph:%0A        - ceph-fuse:%0A        - pexec:%0A            clients:%0A              - ln -s /tmp/cephtest/mnt.* /tmp/cephtest/gmnt%0A        - ssh_keys:%0A        - mpi:%0A            exec: fsx-mpi%0A            workdir: /tmp/cephtest/gmnt%0A        - pexec:%0A            clients:%0A              - rm -f /tmp/cephtest/gmnt%0A%0A
     %22%22%22%0A    
@@ -2681,16 +2681,110 @@
 emote)%0A%0A
+    workdir = %5B%5D%0A    if 'workdir' in config:%0A        workdir = %5B'-wdir', config%5B'workdir'%5D %5D%0A%0A
     log.
@@ -3220,16 +3220,41 @@
 hosts'%5D%0A
+    args.extend(workdir)%0A
     args

Allow working directory to be specified