Python mpi4py.MPI.ANY_TAG Examples

def worker():
    from mpi4py import MPI
    while True:
        status = MPI.Status()
        ret = MPI.COMM_WORLD.recv(source=0, tag=MPI.ANY_TAG, status=status) 

        if status.tag == 10:
            # Workload received
            func = ret['func']
            if ret.get('unpack'):
                res = func(*ret['input_data'])
            else:
                res = func(ret['input_data'])

            # Done, let's send it back
            MPI.COMM_WORLD.send(dict(job_index=ret['job_index'], output_data=res), dest=0, tag=2)

        elif status.tag == 666:
            # Kill code
            sys.exit(0) 

def _read(self):
        result = super(MPIWrapper, self)._read()
        if result is not None:
            return result

        status = MPI.Status()
        msg = comm.recv(source=MPI.ANY_SOURCE, tag=MPI.ANY_TAG, status=status)
        tag = status.Get_tag()
        while tag == STATUS_TERMINATED:
            self.terminated += 1
            if self.terminated >= self._num_sources:
                break
            else:
                msg = comm.recv(source=MPI.ANY_SOURCE,
                                tag=MPI.ANY_TAG,
                                status=status)
                tag = status.Get_tag()
        return msg, tag 

def worker_process(comm, rank, tags, status):
    # Worker processes execute code below
    name = MPI.Get_processor_name()
    print("I am a worker with rank %d on %s." % (rank, name))
    comm.send(None, dest=0, tag=tags.READY)

    while True:
        print('Recieving ...')
        task = comm.recv(source=0, tag=MPI.ANY_TAG, status=status)
        print('received!')

        tag = status.Get_tag()

        if tag == tags.START:
            # Do the work here
            result = task + 1
            print('attempting to send ...')
            comm.send(result, dest=0, tag=tags.DONE)
            print('sending worked ...')
        elif tag == tags.EXIT:
            print('went through exit')
            break 

def slave_set(self):
        if self.rank > 0:
#            print("SLAVE : ", self.rank, " starting...")
            status = MPI.Status()
#            print("SLAVE : ", self.rank, " probing for message...")
            msg = self.COMM.Probe(0, MPI.ANY_TAG, status=status)
#            print("SLAVE : ", self.rank, " recieved a message... ", status.Get_tag())
            self.working = True
            if status.Get_tag() == tags.WORK:
                workingBlock = self.COMM.recv(source=0, tag=tags.WORK)
#                print("SLAVE : ", self.rank, " just recieved ", workingBlock)
                self.curr_block = workingBlock
                self.working = True
                return True, workingBlock
            else:
                self.working = False
                workingBlock = self.COMM.recv(source=0, tag=tags.KILL)
#                print("SLAVE : ", self.rank, " dying...")
                return False, 0
        else:
            return False, 0 

def Recv(self, buf, source=MPI.ANY_SOURCE, tag=MPI.ANY_TAG, status=None):
        if isinstance(buf, dndarray.DNDarray):
            buf = buf._DNDarray__array
        if not isinstance(buf, torch.Tensor):
            return self.handle.Recv(buf, source, tag, status)

        rbuf = buf if CUDA_AWARE_MPI else buf.cpu()
        ret = self.handle.Recv(self.as_buffer(rbuf), source, tag, status)

        if isinstance(buf, torch.Tensor) and buf.is_cuda and not CUDA_AWARE_MPI:
            buf.copy_(rbuf)
        return ret 

def imap(f, workloads, star=False):
    global _g_available_workers, _g_initialized
    from mpi4py import MPI
    N = MPI.COMM_WORLD.Get_size() - 1
    if N == 0 or not _g_initialized:
        mapf = [map, itr.starmap][star]
        for res in mapf(f, workloads):
            yield res
        return

    results = []
    indices = []

    for job_index, workload in enumerate(itr.chain(workloads, itr.repeat(None))):
        if workload is None and len(_g_available_workers) == N:
            break

        while not _g_available_workers or workload is None:
            # Wait to receive results
            status = MPI.Status()
            ret = MPI.COMM_WORLD.recv(source=MPI.ANY_SOURCE, tag=MPI.ANY_TAG, status=status)
            if status.tag == 2:
                results.append(ret['output_data'])
                indices.append(ret['job_index'])
                _g_available_workers.add(status.source)
                if len(_g_available_workers) == N:
                    break

        if _g_available_workers and workload is not None:
            dest_rank = _g_available_workers.pop()

            # Send off job
            task = dict(func=f, input_data=workload, job_index=job_index, unpack=star)
            MPI.COMM_WORLD.send(task, dest=dest_rank, tag=10)

    II = np.argsort(indices)  
    for i in II:
        yield results[i] 

def imap_unordered(f, workloads, star=False):
    global _g_available_workers, _g_initialized

    from mpi4py import MPI
    N = MPI.COMM_WORLD.Get_size() - 1
    if N == 0 or not _g_initialized:
        mapf = [map, itr.starmap][star]
        for res in mapf(f, workloads):
            yield res
        return

    for job_index, workload in enumerate(itr.chain(workloads, itr.repeat(None))):
        if workload is None and len(_g_available_workers) == N:
            break

        while not _g_available_workers or workload is None:
            # Wait to receive results
            status = MPI.Status()
            ret = MPI.COMM_WORLD.recv(source=MPI.ANY_SOURCE, tag=MPI.ANY_TAG, status=status)
            if status.tag == 2:
                yield ret['output_data']
                _g_available_workers.add(status.source)
                if len(_g_available_workers) == N:
                    break

        if _g_available_workers and workload is not None:
            dest_rank = _g_available_workers.pop()

            # Send off job
            task = dict(func=f, input_data=workload, job_index=job_index, unpack=star)
            MPI.COMM_WORLD.send(task, dest=dest_rank, tag=10) 

def recv(self, obj=None, tag=MPI.ANY_TAG, source=None, buffer=False, status=None, comm=None):
        """Wrapper around MPI.recv/Recv. Returns the received object.
            Params:
              obj: variable into which the received object should be placed
              tag: string indicating which MPI tag should be received
              source: integer rank of the message source.  Defaults to self.parent_rank
              buffer: True if the received object should be sent as a single-segment buffer
                (e.g. for numpy arrays) using MPI.Recv rather than MPI.recv
              status: MPI status object that is filled with received status information
              comm: MPI communicator to use.  Defaults to self.parent_comm"""
        if comm is None:
            comm = self.parent_comm
        if source is None:
            if self.parent_rank is None:
                raise Error("Attempting to receive %s from parent, but parent rank is None" % tag)
            source = self.parent_rank 
        tag_num = self.lookup_mpi_tag(tag)
        if tag == 'history':
            obj = comm.recv( source=source, tag=tag_num, status=status )
            return obj
        #if tag in ['bool','time']:
        #    comm.Recv(obj, source=source, tag=tag_num, status=status )
        #    return obj
        if buffer:
            if type(obj) == list:
                for o in obj:
                    if type(o) == list:
                        for sub_o in o:
                            comm.Recv( sub_o, source=source, tag=tag_num, status=status )
                    else:
                        comm.Recv( o, source=source, tag=tag_num, status=status )
            else:
                comm.Recv( obj, source=source, tag=tag_num, status=status )
            return obj
        else:
            obj = comm.recv( source=source, tag=tag_num, status=status )
            return obj 

def manage(comm, num_workers):
    available_to_pair = -1
    finished = {}
    num_finished = 0
    
    for worker_id in xrange(1, num_workers + 1):
        finished[worker_id] = False

    while num_finished != num_workers:
        status = MPI.Status()
        worker_id = comm.recv(source=MPI.ANY_SOURCE, tag=MPI.ANY_TAG, \
                status=status)
        event = status.Get_tag()

        if event == Msg.STARTED.value:
            print('Worker', worker_id, 'is working!')
        
        elif event == Msg.PAIRME.value:
            if num_finished == num_workers - 1: #only 1 working, pair with self
                comm.isend(worker_id, dest=worker_id, tag=Msg.PAIRED.value)
            else:
                assert available_to_pair != worker_id
                if available_to_pair == -1:
                    available_to_pair = worker_id
                else:
                    comm.isend(available_to_pair, dest=worker_id, \
                            tag=Msg.PAIRED.value)
                    comm.isend(worker_id, dest=available_to_pair, \
                            tag=Msg.PAIRED.value)
                    available_to_pair = -1
        elif event == Msg.FINISHED.value:
            print('Worker', worker_id, 'has finished it\'s iterations!')
            finished[worker_id] = True
            num_finished += 1
            
            #wake up last worker if it's waiting for a pair
            if num_finished == num_workers - 1 and available_to_pair != -1:
                comm.isend(available_to_pair, dest=available_to_pair, \
                        tag=Msg.PAIRED.value)
        else:
            print(0, 'Unknown message received', worker_id, event, Msg(event)) 

def work():
    comm = MPI.COMM_WORLD
    rank = comm.rank
    
    #pr = cProfile.Profile()
    #pr.enable()

    while True:
        status = MPI.Status()
        msg = comm.recv(source=MASTER, tag=MPI.ANY_TAG, status=status)
        event = status.Get_tag()

        if event == Msg.LEARN.value:
            comm.isend(rank, dest=MASTER, tag=Msg.STARTED.value)

            num_iter = msg

            Dts, Trace, Count_zh, Count_sz, count_h, count_z, \
                    alpha_zh, beta_zs, kernel = receive_workload(comm)
            fast_populate(Trace, Count_zh, Count_sz, count_h, \
                    count_z)
            sample(Dts, Trace, Count_zh, Count_sz, count_h, \
                    count_z, alpha_zh, beta_zs, kernel, num_iter, \
                    comm)
            
            comm.isend(rank, dest=MASTER, tag=Msg.FINISHED.value)
        elif event == Msg.SENDRESULTS.value:
            comm.Send([np.array(Trace[:, -1], order='C'), MPI.INT], dest=MASTER)
            comm.Send([Count_zh, MPI.INT], dest=MASTER)
            comm.Send([Count_sz, MPI.INT], dest=MASTER)
            comm.Send([count_h, MPI.INT], dest=MASTER)
            comm.Send([count_z, MPI.INT], dest=MASTER)
            comm.Send([kernel.get_state(), MPI.DOUBLE], dest=MASTER)
        elif event == Msg.STOP.value:
            break
        else:
            print('Unknown message received', msg, event, Msg(event))

    #pr.disable()
    #pr.dump_stats('worker-%d.pstats' % rank) 

def wait_for_root(self):
        """
        Wait for message from root. Normally these are the updated parameters.
        @return The sent data from the root.
        """
        tensor = self.comm.recv(source=0, tag=MPI.ANY_TAG)
        return tensor 

def wait_for_any_rank(self):
        """
        Waits for any message from any rank.
        @return A 3-tuple of (tensor, source rank, tag)
        """
        status = MPI.Status()
        tensor = self.comm.recv(source=MPI.ANY_SOURCE, tag=MPI.ANY_TAG, status=status)
        return tensor, status.source, status.tag 

def worker(self):
                '''worker processor method which waits for map data or end signal from master'''
                status = MPI.Status()
                while True:
                        job = self.comm.recv(source=0, tag=MPI.ANY_TAG, status=status)
                        if job == -999: #end signal
                                break
                        if  isinstance(job, _func_wrapper):
                                self.function = job.function
                                continue

                        result = self.function(job)
                        self.comm.isend(result, dest=0, tag=status.tag) 

def Irecv(self, buf, source=MPI.ANY_SOURCE, tag=MPI.ANY_TAG):
        if isinstance(buf, dndarray.DNDarray):
            buf = buf._DNDarray__array
        if not isinstance(buf, torch.Tensor):
            return MPIRequest(self.handle.Irecv(buf, source, tag))

        rbuf = buf if CUDA_AWARE_MPI else buf.cpu()
        return MPIRequest(self.handle.Irecv(self.as_buffer(rbuf), source, tag), None, rbuf, buf) 

def _worker(self, clf):
        """Worker node's operation.

        Receiving tasks from the master to process and sending the result back

        Parameters
        ----------
        clf: classification function
            the classifier to be used in cross validation

        Returns
        -------
        None
        """
        logger.debug(
            'worker %d is running, waiting for tasks from master at rank %d' %
            (MPI.COMM_WORLD.Get_rank(), self.master_rank)
        )
        comm = MPI.COMM_WORLD
        status = MPI.Status()
        while 1:
            task = comm.recv(source=self.master_rank,
                             tag=MPI.ANY_TAG,
                             status=status)
            if status.Get_tag():
                break
            comm.send(self._voxel_scoring(task, clf),
                      dest=self.master_rank) 

def receive(wrapper):
    while wrapper.terminated < wrapper._num_sources:
        status = MPI.Status()
        msg = comm.recv(source=MPI.ANY_SOURCE, tag=MPI.ANY_TAG, status=status)
        tag = status.Get_tag()
        # print('Received %s, %s' % (msg, tag))
        if tag == STATUS_TERMINATED:
            wrapper.terminated += 1
        else:
            wrapper.input_data.put((msg, tag))
        # self.wrapper.pe.log('Queue size: %s'%self.wrapper.input_data.qsize())
    # put the final terminate block into the queue
    wrapper.input_data.put((None, STATUS_TERMINATED)) 

def _create_worker(self):
        comm = self.comm

        while True:

            status = MPI.Status()
            comm.Probe(self.master_rank, MPI.ANY_TAG, status)
            message_src = self.master_rank
            message_tag = status.Get_tag()

            # Check for available task 
            if message_tag == self.task_tag:

                task = comm.recv(source = message_src, tag = message_tag)

                try:
                    result_value = task.fn(*task.args, **task.kwargs)
                except Exception as e:
                    result_object = (task.task_id, 'exception', result_value)
                else:
                    result_object = (task.task_id, 'result', result_value)

                comm.send(result_object, dest = self.master_rank, tag = self.result_tag)

            # Check for announcements
            if message_tag == self.announce_tag:
                messages = comm.recv(source = message_src, tag = message_tag)
                if 'shutdown' in messages:
                    return 

def _receive_loop(self):
        comm = self.comm 

        while True:

            status = MPI.Status()
            comm.Iprobe(MPI.ANY_SOURCE, MPI.ANY_TAG, status)
            message_src = status.Get_source()
            message_tag = status.Get_tag()

            # results are tuples of (task_id, {'result', 'exception'}, value)
            if message_tag == self.result_tag:
                (task_id, result_stat, result_value) = comm.recv(source = message_src, tag = message_tag)

                ft = self.pending_futures.pop(task_id)

                if result_stat == 'exception':
                    ft._set_exception(*result_value)
# Check with Matt on what else to do for an exception
                else:
                    ft._set_result(result_value)
                    self.task_dest.append(message_src)

            # Check for announcements
            elif message_tag == self.announce_tag:
                messages = comm.recv(source = message_src, tag = message_tag)
                if 'shutdown' in messages:
                    log.debug('exiting _receive_loop()')
                    return 

def wait(self):
        '''
        If this isn't the master process, wait for instructions.

        '''
        if self.is_master():
            raise RuntimeError('Master node told to await jobs.')

        status = MPI.Status()

        while True:
            # Event loop.
            # Sit here and await instructions.
            if self.debug:
                print('Worker {0} waiting for task.'.format(self.rank))

            # Blocking receive to wait for instructions.
            task = self.comm.recv(source=0, tag=MPI.ANY_TAG, status=status)
            if self.debug:
                print('Worker {0} got task {1} with tag {2}.'
                      .format(self.rank, task, status.tag))

            # Check if message is special sentinel signaling end.
            # If so, stop.
            if isinstance(task, _close_pool_message):
                if self.debug:
                    print('Worker {0} told to quit.'.format(self.rank))
                break

            # Check if message is special type containing new function
            # to be applied
            if isinstance(task, _function_wrapper):
                self.function = task.function
                if self.debug:
                    print('Worker {0} replaced its task function: {1}.'
                          .format(self.rank, self.function))
                continue

            # If not a special message, just run the known function on
            # the input and return it asynchronously.
            result = self.function(task)
            if self.debug:
                print('Worker {0} sending answer {1} with tag {2}.'
                      .format(self.rank, result, status.tag))
            self.comm.isend(result, dest=0, tag=status.tag) 

def wait(self):
        """
        If this isn't the master process, wait for instructions.

        """
        if self.is_master():
            raise RuntimeError("Master node told to await jobs.")

        status = MPI.Status()

        while True:
            # Event loop.
            # Sit here and await instructions.
            if self.debug:
                print("Worker {0} waiting for task.".format(self.rank))

            # Blocking receive to wait for instructions.
            task = self.comm.recv(source=0, tag=MPI.ANY_TAG, status=status)

            if self.debug:
                print("Worker {0} got task {1} with tag {2}."
                      .format(self.rank, type(task), status.tag))

            # Check if message is special sentinel signaling end.
            # If so, stop.
            if isinstance(task, _close_pool_message):
                if self.debug:
                    print("Worker {0} told to quit.".format(self.rank))
                break

            # Check if message is special type containing new function
            # to be applied
            if isinstance(task, _function_wrapper):
                self.function = task.function
                if self.debug:
                    print("Worker {0} replaced its task function: {1}."
                          .format(self.rank, self.function))
                continue

            # If not a special message, just run the known function on
            # the input and return it asynchronously.
            result = self.function(task)
            if self.debug:
                print("Worker {0} sending answer {1} with tag {2}."
                      .format(self.rank, type(result), status.tag))
            self.comm.isend(result, dest=0, tag=status.tag)

        # Kill the process?
        if self.exit_on_end:
            sys.exit() 

def map(self, function, jobs):
                '''
                Map function to perform similar function to multiprocessing mp.pool map()
                Input:
                function  - function to be mapped
                jobs - array of jobs to be assigned to each proc

                '''
                njobs = len(jobs)
                self.function = function


                # If not the master just wait for instructions.
                if not self.rank == 0:
                        self.worker()
                        return

                F = _func_wrapper(function)
                req = [self.comm.isend(F, dest=i) for i in range(1,self.size)]
                MPI.Request.waitall(req)

                if  (njobs <= self.size-1):
                        requests = []
                        for i, job in enumerate(jobs):
                                worker_id = i % self.size + 1
                                req = self.comm.isend(job, dest=worker_id, tag=i)
                                requests.append(req)

                        MPI.Request.waitall(requests)

                        results = []
                        for i in range(njobs):
                                worker_id = i % self.size + 1
                                result = self.comm.recv(source=worker_id, tag=i)
                                results.append(result)
                        return results

                else:
                        for i in range(1, self.size):
                                req = self.comm.isend(jobs[i-1], dest=i, tag=i)

                        njobs_done = self.size-1 # don't use master for function
                        results = [None]*njobs
                        for job in range(njobs):
                                status = MPI.Status()
                                result = self.comm.recv(source=MPI.ANY_SOURCE,tag=MPI.ANY_TAG, status=status)
                                worker_id = status.source
                                results[job] = result

                                if njobs_done < njobs:
                                        job = jobs[njobs_done]
                                        i = njobs_done
                                        self.comm.isend(job, dest=worker_id, tag=i)
                                        njobs_done += 1


                        return results 

def master_process(comm, size, tags, status):

    num_workers = size - 1
    tasks = range(num_workers)
    chain = []
    active_workers = 0
    # start sampling of chains with given seed
    print("Master starting with %d workers" % num_workers)
    for i in range(num_workers):
        comm.recv(source=MPI.ANY_SOURCE, tag=tags.READY, status=status)
        source = status.Get_source()
        comm.send(tasks[i], dest=source, tag=tags.START)
        print("Sent task to worker %i" % source)
        active_workers += 1

    print("Parallel tempering ...")
    print("----------------------")

    while True:
        m1 = comm.recv(source=MPI.ANY_SOURCE, tag=MPI.ANY_TAG, status=status)
        source1 = status.Get_source()
        print("Got sample 1 from worker %i" % source1)
        m2 = comm.recv(source=MPI.ANY_SOURCE, tag=MPI.ANY_TAG, status=status)
        source2 = status.Get_source()
        print("Got sample 2 from worker %i" % source1)

        m1, m2 = metrop_select(m1, m2)
        print('samples 1, 2 %i %i' % (m1, m2))
        chain.extend([m1, m2])
        if len(chain) < nsamples:
            print("Sending states back to workers ...")
            comm.send(m1, dest=source1, tag=tags.START)
            comm.send(m2, dest=source2, tag=tags.START)
        else:
            print('Requested number of samples reached!')
            break

    print("Master finishing, recorded chain:")
    print(chain)
    print("Closing ...")
    for i in range(1, size):
        print('sending signal to close to %i' % i)
        comm.send(None, dest=i, tag=tags.EXIT)

        print("Closed worker %i" % i)
        active_workers -= 1 

def worker_process(comm, tags, status):
    """
    Worker processes, that do the actual sampling.
    They receive all arguments by the master process.

    Parameters
    ----------
    comm : mpi.communicator
    tags : message tags
    status : mpi.status object
    """
    name = MPI.Get_processor_name()
    logger.debug(
        "Entering worker process with rank %d on %s." % (comm.rank, name))
    comm.send(None, dest=0, tag=tags.READY)

    logger.debug('Worker %i receiving work package ...' % comm.rank)
    kwargs = comm.recv(source=0, tag=tags.INIT, status=status)
    logger.debug('Worker %i received package!' % comm.rank)

    try:
        step = kwargs['step']
    except KeyError:
        raise ValueError('Step method not defined!')

    # do initial sampling
    result = sample_pt_chain(**kwargs)
    comm.Send([result, MPI.DOUBLE], dest=0, tag=tags.DONE)

    # enter repeated sampling
    while True:
        # TODO: make transd-compatible
        data = num.empty(step.lordering.size, dtype=tconfig.floatX)
        comm.Recv([data, MPI.DOUBLE],
                  tag=MPI.ANY_TAG, source=0, status=status)

        tag = status.Get_tag()
        if tag == tags.SAMPLE:
            lpoint = step.lij.a2l(data)
            start = step.lij.l2d(lpoint)
            kwargs['start'] = start
            # overwrite previous point in case got swapped
            kwargs['step'].chain_previous_lpoint[comm.rank] = lpoint
            result = sample_pt_chain(**kwargs)

            logger.debug('Worker %i attempting to send ...' % comm.rank)
            comm.Send([result, MPI.DOUBLE], dest=0, tag=tags.DONE)
            logger.debug('Worker %i sent message successfully ...' % comm.rank)

        elif tag == tags.BETA:
            logger.debug(
                'Worker %i received beta: %f' % (comm.rank, data[0]))
            kwargs['step'].beta = data[0]

        elif tag == tags.EXIT:
            logger.debug('Worker %i went through EXIT!' % comm.rank)
            break 

def orchestrate_map(self,pds_id):
        """Orchestrates the teams to perform a map function
        
        This works by keeping track of the teams who haven't finished executing,
        waiting for them to request the next chunk of data when they are free,
        responding to them with the data and then sending them a Sentinel
        signalling that they can exit.
        """
        is_map_done = [True if i in self.mpimanager.get_scheduler_node_ranks() else False for i in range(self.mpimanager.get_scheduler_size())]
        status = MPI.Status()

        #Copy it to the pending. This is so when scheduler accesses
        #the PDS data it's not empty.
        self.pds_pending_store[pds_id] = list(self.pds_store[pds_id])

        #While we have some ranks that haven't finished
        while sum(is_map_done)<self.mpimanager.get_scheduler_size():
            #Wait for a reqest from anyone
            data_request = self.mpimanager.get_scheduler_communicator().recv(
                source=MPI.ANY_SOURCE,
                tag=MPI.ANY_TAG,
                status=status,
            )
            request_from_rank = status.source

            if data_request!=pds_id:
                print("Ignoring stale PDS data request from",
                    request_from_rank,":",data_request,"/",pds_id)
                continue

            #Pointer so we don't have to keep doing dict lookups
            current_pds_items = self.pds_pending_store[pds_id]
            num_current_pds_items = len(current_pds_items)

            #Everyone's already exhausted all the data.
            # Send a sentinel and mark the node as finished
            if num_current_pds_items == 0:
                self.mpimanager.get_scheduler_communicator().send(None, dest=request_from_rank, tag=pds_id)
                is_map_done[request_from_rank] = True
            else:
                #Create the chunk of data to send. Pop off items and tag them with an id.
                # so we can sort them later
                chunk_to_send = []
                for i in range(self.chunk_size):
                    chunk_to_send+=[(num_current_pds_items-i,current_pds_items.pop())]
                    self.mpimanager.get_scheduler_communicator().send(chunk_to_send, dest=request_from_rank, tag=pds_id) 

def _distribute_tasks(self, tasks):
        """
        Internal function that distributes the tasks from the root to the workers
        """
        if not self.is_root():
            raise ValueError("only the root rank should distribute the tasks")

        ntasks = len(tasks)
        task_index     = 0
        closed_workers = 0

        # logging info
        args = (self.workers, ntasks)
        self.logger.debug("master starting with %d worker(s) with %d total tasks" %args)

        # loop until all workers have finished with no more tasks
        while closed_workers < self.workers:

            # look for tags from the workers
            data = self.basecomm.recv(source=MPI.ANY_SOURCE, tag=MPI.ANY_TAG, status=self.status)
            source = self.status.Get_source()
            tag = self.status.Get_tag()

            # worker is ready, so send it a task
            if tag == self.tags.READY:

                # still more tasks to compute
                if task_index < ntasks:
                    this_task = [task_index, tasks[task_index]]
                    self.basecomm.send(this_task, dest=source, tag=self.tags.START)
                    self.logger.debug("sending task `%s` to worker %d" %(str(tasks[task_index]), source))
                    task_index += 1

                # all tasks sent -- tell worker to exit
                else:
                    self.basecomm.send(None, dest=source, tag=self.tags.EXIT)

            # store the results from finished tasks
            elif tag == self.tags.DONE:
                self.logger.debug("received result from worker %d" %source)

            # track workers that exited
            elif tag == self.tags.EXIT:
                closed_workers += 1
                self.logger.debug("worker %d has exited, closed workers = %d" %(source, closed_workers)) 

def _get_tasks(self):
        """
        Internal generator that yields the next available task from a worker
        """
        if self.is_root():
            raise RuntimeError("Root rank mistakenly told to await tasks")

        # logging info
        if self.comm.rank == 0:
            args = (self.rank, MPI.Get_processor_name(), self.comm.size)
            self.logger.debug("worker master rank is %d on %s with %d processes available" %args)

        # continously loop and wait for instructions
        while True:
            args = None
            tag = -1

            # have the master rank of the subcomm ask for task and then broadcast
            if self.comm.rank == 0:
                self.basecomm.send(None, dest=0, tag=self.tags.READY)
                args = self.basecomm.recv(source=0, tag=MPI.ANY_TAG, status=self.status)
                tag = self.status.Get_tag()

            # bcast to everyone in the worker subcomm
            args  = self.comm.bcast(args) # args is [task_number, task_value]
            tag   = self.comm.bcast(tag)

            # yield the task
            if tag == self.tags.START:

                # yield the task value
                yield args

                # wait for everyone in task group before telling master this task is done
                self.comm.Barrier()
                if self.comm.rank == 0:
                    self.basecomm.send([args[0], None], dest=0, tag=self.tags.DONE)

            # see ya later
            elif tag == self.tags.EXIT:
                break

        # wait for everyone in task group and exit
        self.comm.Barrier()
        if self.comm.rank == 0:
            self.basecomm.send(None, dest=0, tag=self.tags.EXIT)

        # debug logging
        self.logger.debug("rank %d process is done waiting" %self.rank) 

def wait(self):
        """
        If this isn't the master process, wait for instructions.

        """
        from mpi4py import MPI
        if self.is_master():
            raise RuntimeError("Master node told to await jobs.")

        status = MPI.Status()

        while True:
            # Event loop.
            # Sit here and await instructions.
            if self.debug:
                print("Worker {0} waiting for task.".format(self.rank))

            # Blocking receive to wait for instructions.
            task = self.comm.recv(source=0, tag=MPI.ANY_TAG, status=status)
            if self.debug:
                print("Worker {0} got task {1} with tag {2}."
                      .format(self.rank, task, status.tag))

            # Check if message is special sentinel signaling end.
            # If so, stop.
            if isinstance(task, _close_pool_message):
                if self.debug:
                    print("Worker {0} told to quit.".format(self.rank))
                break

            # Check if message is special type containing new function
            # to be applied
            if isinstance(task, _function_wrapper):
                self.function = task.function
                if self.debug:
                    print("Worker {0} replaced its task function: {1}."
                          .format(self.rank, self.function))
                continue

            # If not a special message, just run the known function on
            # the input and return it asynchronously.
            result = self.function(task)
            if self.debug:
                print("Worker {0} sending answer {1} with tag {2}."
                      .format(self.rank, result, status.tag))
            self.comm.isend(result, dest=0, tag=status.tag) 

def wait(self):
            """
            If this isn't the master process, wait for instructions.

            """
            if self.is_master():
                raise RuntimeError("Master node told to await jobs.")

            status = MPI.Status()

            while True:
                # Event loop.
                # Sit here and await instructions.
                if self.debug:
                    print("Worker {0} waiting for task.".format(self.rank))

                # Blocking receive to wait for instructions.
                task = self.comm.recv(source=0, tag=MPI.ANY_TAG, status=status)
                if self.debug:
                    print("Worker {0} got task {1} with tag {2}."
                                     .format(self.rank, task, status.tag))

                # Check if message is special sentinel signaling end.
                # If so, stop.
                if isinstance(task, _close_pool_message):
                    if self.debug:
                        print("Worker {0} told to quit.".format(self.rank))
                    break

                # Check if message is special type containing new function
                # to be applied
                if isinstance(task, _function_wrapper):
                    self.function = task.function
                    if self.debug:
                        print("Worker {0} replaced its task function: {1}."
                                .format(self.rank, self.function))
                    continue

                # If not a special message, just run the known function on
                # the input and return it asynchronously.
                result = self.function(task)
                if self.debug:
                    print("Worker {0} sending answer {1} with tag {2}."
                            .format(self.rank, result, status.tag))
                self.comm.isend(result, dest=0, tag=status.tag)