Python mpi4py.MPI.Status() Examples

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 _handle_message(self):
        """Handle received messages.

        Receive record update messages of the form
            ('action', record_id, params)
        where 'action' is the name of an EvalRecord method and params is
        the list of parameters.  The record_id should be recorded in the
        hub's records table (which happens whenever it is referenced in
        a message sent to a worker).

        On a message indicating that the worker is done with the record,
        we add the worker that sent the message back to the free pool.
        """
        logger.debug("Handle incoming message")
        s = MPI.Status()
        data = comm.recv(status=s, source=MPI.ANY_SOURCE, tag=0)
        logger.debug("Received message: %s", data)
        mname = data[0]
        record = self._recids[data[1]]
        method = getattr(record, mname)
        method(*data[2:])
        if mname == 'complete' or mname == 'cancel' or mname == 'kill':
            logger.debug("Re-queueing worker")
            self._workers.append(s.source)
        self.ping() 

def _dispatch_loop(self):
        comm = self.comm        

        while True:

            # Dispatch as many tasks as possible before checking for shutdown
            while self.task_dest:
                try:
                    task = self.task_queue.popleft()
                    task_dest = self.task_dest.popleft()
                except IndexError:
                    break
                else:
                    comm.send(task, dest = task_dest, tag = self.task_tag )

            status = MPI.Status()
            comm.Iprobe(self.master_rank, self.announce_tag, status)
            message_tag = status.Get_tag()

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

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 _write(self, name, data):
        try:
            targets = self.targets[name]
        except KeyError:
            # no targets
            # self.pe.log('Produced output: %s' % {name: data})
            return
        for (inputName, communication) in targets:
            output = {inputName: data}
            dest = communication.getDestination(output)
            for i in dest:
                try:
                    # self.pe.log('Sending %s to %s' % (output, i))
                    request = comm.isend(output, tag=STATUS_ACTIVE, dest=i)
                    status = MPI.Status()
                    request.Wait(status)
                except:
                    self.pe.log(
                        'Failed to send data stream "%s" to rank %s: %s'
                        % (name, i, traceback.format_exc())) 

def __init__(self, rank=None, nonMPIMode=False, logger=None):
        if nonMPIMode:
            self.comm = None
            self.stat = None
            self.nRank = 0
            self.totalRanks = 1
            self.selectSource = None
        else:
            from mpi4py import MPI
            self.comm = MPI.COMM_WORLD
            self.stat = MPI.Status()
            self.nRank = self.comm.Get_rank()
            self.totalRanks = self.comm.Get_size()
            self.selectSource = MPI.ANY_SOURCE

        self.logger = logger

        # for message in rank 0
        self.hasMessage = False
        self.recvQueue = recvQueue
        self.sendQueue = sendQueue


    # get rank of itself 

def mpi_master_spin_tasks(task, gdirs):
    comm = OGGM_MPI_COMM
    cfg_store = cfg.pack_config()
    msg_list = ([gdir for gdir in gdirs if gdir is not None] +
                [None] * OGGM_MPI_SIZE)

    _imprint("Starting MPI task distribution...")

    comm.bcast((cfg_store, task), root=OGGM_MPI_ROOT)

    status = MPI.Status()
    for msg in msg_list:
        comm.recv(source=MPI.ANY_SOURCE, status=status)
        comm.send(obj=msg, dest=status.Get_source())

    _imprint("MPI task distribution done, collecting results...")

    comm.gather(sendobj=None, root=OGGM_MPI_ROOT)

    _imprint("MPI task results gotten!") 

def __init__(self, cpus_per_task, comm=None, debug=False, use_all_cpus=False):

        if debug:
            self.logger.setLevel(logging.DEBUG)

        self.cpus_per_task = cpus_per_task
        self.use_all_cpus  = use_all_cpus

        # the base communicator
        self.basecomm = MPI.COMM_WORLD if comm is None else comm
        self.rank     = self.basecomm.rank
        self.size     = self.basecomm.size

        # need at least one
        if self.size == 1:
            raise ValueError("need at least two processes to use a TaskManager")

        # communication tags
        self.tags = enum('READY', 'DONE', 'EXIT', 'START')

        # the task communicator
        self.comm = None

        # store a MPI status
        self.status = MPI.Status() 

def __init__(self, comm, **kwargs):
        super(NN_Trainer, self).__init__()
        self.comm = comm   # get MPI communicator object
        self.world_size = comm.Get_size() # total number of processes
        self.rank = comm.Get_rank() # rank of this Worker
        #self.status = MPI.Status()
        self.cur_step = 0
        self.next_step = 0 # we will fetch this one from parameter server

        self.batch_size = kwargs['batch_size']
        self.max_epochs = kwargs['max_epochs']
        self.momentum = kwargs['momentum']
        self.lr = kwargs['learning_rate']
        self._max_steps = kwargs['max_steps']
        self.network_config = kwargs['network']
        self.comm_type = kwargs['comm_method']
        self.kill_threshold = kwargs['kill_threshold']
        self._eval_batch_size = 100
        self._eval_freq = kwargs['eval_freq']
        self._train_dir = kwargs['train_dir']
        self._compress_grad = kwargs['compress_grad']
        self._device = kwargs['device']

        # this one is going to be used to avoid fetch the weights for multiple times
        self._layer_cur_step = [] 

def pt_sample():
    # Define MPI message tags
    tags = enum('READY', 'DONE', 'EXIT', 'START')

    # Initializations and preliminaries
    comm = MPI.COMM_WORLD   # get MPI communicator object
    size = comm.size        # total number of processes
    rank = comm.rank        # rank of this process
    status = MPI.Status()   # get MPI status object

    if rank == 0:
        print('Here')
        master_process(comm, size, tags, status)
    else:
        print('worker')
        worker_process(comm, rank, tags, status)

    print('Done!') 

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 _handle_message(self):
        "Handle an incoming message and dispatch to handler."
        logger.debug("Handle incoming message")
        try:
            s = MPI.Status()
            data = comm.recv(status=s, source=0, tag=0)
            logger.debug("Incoming message received")
            mname = "on_{0}".format(data[0])
            logger.debug("Call to %s%s", mname, data[1:])
            method = getattr(self, mname)
            method(*data[1:])
        except Exception:
            logger.debug("Exception in message handler") 

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 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 train(self):
        """Broadcasts model information to children and signals them to start training.
            Receive messages from workers and processes each message until training is done.
            When finished, signal the parent process that training is complete.
        """
        Trace.begin("train")
        self.start_time = time.time()
        self.check_sanity()
        self.bcast_weights( comm=self.child_comm )
        self.signal_children()

        if self.threaded_validation:
            self.validation_queue = queue.Queue()
            self.validation_thread = threading.Thread(target=MPIMaster.validation_worker, args = (self, ))
            self.validation_thread.start()

        status = MPI.Status()
        self.running_workers = list(range(1, self.num_workers+1))
        self.waiting_workers_list = []
        
        while self.running_workers:
            self.recv_any_from_child(status)
            self.process_message( status )
            if (self.stop_training):
                self.shut_down_workers()
        self.logger.info("Done training")
        # If we did not finish the last epoch, validate one more time.
        # (this happens if the batch size does not divide the dataset size)
        if self.epoch < self.num_epochs or not self.histories.get(self.history_key(),None):
            epoch_logs = self.validate(self.weights)
        if self.threaded_validation:
            self.validation_queue.put(None)
            self.validation_queue.join()
            self.validation_thread.join()
        self.save_checkpoint()
        self.send_exit_to_parent()
        self.send_history_to_parent()
        self.data.finalize()
        self.stop_time = time.time()
        Trace.end("train") 

def _gather_updated_ages(self, total_age):
        total_age.update({0: self._island.generational_age})
        status = MPI.Status()
        while self.comm.iprobe(source=MPI.ANY_SOURCE,
                               tag=AGE_UPDATE,
                               status=status):
            data = self.comm.recv(source=status.Get_source(),
                                  tag=AGE_UPDATE)
            total_age.update(data) 

def recv_result_from_workers(self):
        """ Receives a results from the MPI worker pool and send it out via 0mq

        Returns:
        --------
            result: task result from the workers
        """
        info = MPI.Status()
        result = self.comm.recv(source=MPI.ANY_SOURCE, tag=RESULT_TAG, status=info)
        logger.debug("Received result from workers: {}".format(result))
        return result 

def recv_task_request_from_workers(self):
        """ Receives 1 task request from MPI comm

        Returns:
        --------
            worker_rank: worker_rank id
        """
        info = MPI.Status()
        comm.recv(source=MPI.ANY_SOURCE, tag=TASK_REQUEST_TAG, status=info)
        worker_rank = info.Get_source()
        logger.info("Received task request from worker:{}".format(worker_rank))
        return worker_rank 

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 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 process_messages(self, count_arr, recorder=None):
        
        if recorder: recorder.start()
        
        status = MPI.Status()
        
        s = self.comm.Iprobe(source=MPI.ANY_SOURCE, tag=700, status=status)
        
        # if self.test: print '%d probed, got %s' % (self.rank,s)
        
        while s:
            
            src_rank=status.source
            
            self.comm.Recv(buf=count_arr, source=src_rank, tag=700, status=status)
            
            self.gpucomm, src_gpurank, self_gpurank = self.get_gpucomm_with(src_rank)
            
            if self.test: print('%d merging with %d' % (self.rank, src_rank))
            
            self._merge_params_from(src_gpurank, src_rank)
            
            s = self.comm.Iprobe(source=MPI.ANY_SOURCE, tag=700, status=status)
            
            if self.test: print('%d probed again, got %s' % (self.rank,s))
            
        if recorder: recorder.end('comm') 

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 __init__(self,comm=None):
                '''initialize a communicator and set the rank and size '''
                if comm==None:
                        self.comm = MPI.COMM_WORLD
                else:
                        self.comm=comm
                self.rank = self.comm.Get_rank()
                self.size = self.comm.Get_size()
                self.status = MPI.Status() 

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 _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 _sample():
    # Define MPI message tags
    tags = distributed.enum('READY', 'INIT', 'DONE', 'EXIT', 'SAMPLE', 'BETA')

    # Initializations and preliminaries
    comm = MPI.COMM_WORLD
    status = MPI.Status()

    logger.debug('communicator size %i' % comm.size)
    model = load_objects(distributed.pymc_model_name)
    with model:
        for i in range(comm.size):
            if i == comm.rank:
                logger.info('Working %i' % i)

        comm.Barrier()

        if comm.rank == 0:
            logger.info('Loading passed arguments ...')

            arguments = load_objects(distributed.mpiargs_name)
            args = [model] + arguments

            master_process(comm, tags, status, *args)
        else:
            worker_process(comm, tags, status) 

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 _shutdown_slaves():
    global OGGM_MPI_COMM
    if OGGM_MPI_COMM is not None and OGGM_MPI_COMM != MPI.COMM_NULL:
        msgs = [StopIteration] * OGGM_MPI_SIZE
        status = MPI.Status()
        OGGM_MPI_COMM.bcast((None, None), root=OGGM_MPI_ROOT)
        for msg in msgs:
            OGGM_MPI_COMM.recv(source=MPI.ANY_SOURCE, status=status)
            OGGM_MPI_COMM.send(obj=msg, dest=status.Get_source())
        OGGM_MPI_COMM.gather(sendobj=None, root=OGGM_MPI_ROOT)
    OGGM_MPI_COMM = None 

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))