# Path: taskloaf/mpi.py
# def mpiexisting(n_workers, f, cfg, die_on_exception=True):
#     orig_sys_except = sys.excepthook
#     if die_on_exception:
# 
#         def die(*args, **kwargs):
#             orig_sys_except(*args, **kwargs)
#             sys.stderr.flush()
#             MPI.COMM_WORLD.Abort()
# 
#         sys.excepthook = die
# 
#     n_mpi_procs = MPI.COMM_WORLD.Get_size()
#     if n_workers > n_mpi_procs:
#         raise Exception(
#             "There are only %s MPI processes but %s were requested"
#             % (n_mpi_procs, n_workers)
#         )
#     c = MPIComm()
#     out = f(c) if c.addr < n_workers else None
#     sys.excepthook = orig_sys_except
#     return out
# 
# class MPIComm:
#     default_comm = MPI.COMM_WORLD
# 
#     # TODO: multi part messages can be done with tags.
#     def __init__(self, comm=None):
#         if comm is None:
#             comm = MPIComm.default_comm
#         self.comm = comm
#         self.addr = rank(self.comm)
# 
#     def send(self, to_addr, data):
#         self.comm.Isend(data, dest=to_addr)
# 
#     def recv(self):
#         s = MPI.Status()
#         msg_exists = self.comm.iprobe(status=s)
#         if not msg_exists:
#             return None
#         out = memoryview(bytearray(s.count))
#         self.comm.Recv(out, source=s.source)
#         return out
# 
#     def barrier(self):
#         self.comm.Barrier()
# 
# def rank(comm=MPI.COMM_WORLD):
#     return comm.Get_rank()
# 
# Path: taskloaf/test_decorators.py
# def mpi_procs(n):
#     n_procs_available = MPI.COMM_WORLD.Get_size()
# 
#     def decorator(test_fnc):
#         test_fnc.n_procs = n
#         return test_fnc
# 
#     return decorator


import os
import pytest
import asyncio
import taskloaf.worker
    import gc; gc.collect()
from taskloaf.cluster import cluster
from taskloaf.mpi import mpiexisting, MPIComm, rank
from taskloaf.test_decorators import mpi_procs
from taskloaf.run import run
from fixtures import w
            taskloaf.worker.shutdown(w)
        w.submit_work(1, g)
        while True:
            await asyncio.sleep(0)
    cluster(2, f, runner = mpiexisting)

def test_cluster_output():
    async def f(w):
        return 1
    assert(cluster(1, f) == 1)

class FunnyException(Exception):
    pass

@mpi_procs(2)
def test_cluster_death_cleansup():
    def check(n, onoff):
        for i in range(n):
            path = '/dev/shm/taskloaf_' + str(i) + '_0'
            assert(os.path.exists(path) == onoff)
    async def f(w):
        ptr = w.allocator.malloc(1)
        check(1, True)
        raise FunnyException()
    check(2, False)
    raises = False
    try:
        cluster(2, f, runner = mpiexisting)
    except FunnyException as e:
        raises = True

# Path: taskloaf/mpi.py
# def mpiexisting(n_workers, f, cfg, die_on_exception=True):
#     orig_sys_except = sys.excepthook
#     if die_on_exception:
# 
#         def die(*args, **kwargs):
#             orig_sys_except(*args, **kwargs)
#             sys.stderr.flush()
#             MPI.COMM_WORLD.Abort()
# 
#         sys.excepthook = die
# 
#     n_mpi_procs = MPI.COMM_WORLD.Get_size()
#     if n_workers > n_mpi_procs:
#         raise Exception(
#             "There are only %s MPI processes but %s were requested"
#             % (n_mpi_procs, n_workers)
#         )
#     c = MPIComm()
#     out = f(c) if c.addr < n_workers else None
#     sys.excepthook = orig_sys_except
#     return out
# 
# class MPIComm:
#     default_comm = MPI.COMM_WORLD
# 
#     # TODO: multi part messages can be done with tags.
#     def __init__(self, comm=None):
#         if comm is None:
#             comm = MPIComm.default_comm
#         self.comm = comm
#         self.addr = rank(self.comm)
# 
#     def send(self, to_addr, data):
#         self.comm.Isend(data, dest=to_addr)
# 
#     def recv(self):
#         s = MPI.Status()
#         msg_exists = self.comm.iprobe(status=s)
#         if not msg_exists:
#             return None
#         out = memoryview(bytearray(s.count))
#         self.comm.Recv(out, source=s.source)
#         return out
# 
#     def barrier(self):
#         self.comm.Barrier()
# 
# def rank(comm=MPI.COMM_WORLD):
#     return comm.Get_rank()
# 
# Path: taskloaf/test_decorators.py
# def mpi_procs(n):
#     n_procs_available = MPI.COMM_WORLD.Get_size()
# 
#     def decorator(test_fnc):
#         test_fnc.n_procs = n
#         return test_fnc
# 
#     return decorator


# Path: tests/test_worker.py
import os
import pytest
import asyncio
import taskloaf.worker
    import gc; gc.collect()
from taskloaf.cluster import cluster
from taskloaf.mpi import mpiexisting, MPIComm, rank
from taskloaf.test_decorators import mpi_procs
from taskloaf.run import run
from fixtures import w

            taskloaf.worker.shutdown(w)
        w.submit_work(1, g)
        while True:
            await asyncio.sleep(0)
    cluster(2, f, runner = mpiexisting)

def test_cluster_output():
    async def f(w):
        return 1
    assert(cluster(1, f) == 1)

class FunnyException(Exception):
    pass

@mpi_procs(2)
def test_cluster_death_cleansup():
    def check(n, onoff):
        for i in range(n):
            path = '/dev/shm/taskloaf_' + str(i) + '_0'
            assert(os.path.exists(path) == onoff)
    async def f(w):
        ptr = w.allocator.malloc(1)
        check(1, True)
        raise FunnyException()
    check(2, False)
    raises = False
    try:
        cluster(2, f, runner = mpiexisting)
    except FunnyException as e:
        raises = True

    assert(rank() == 1 or raises)