Class OrderedResultGatherProc

tensorpack/utils/concurrency.py:318–358

Gather indexed data from a data queue, and produce results with the original index-based order.

class OrderedResultGatherProc(mp.Process):
    """
    Gather indexed data from a data queue, and produce results with the
    original index-based order.
    """

    def __init__(self, data_queue, nr_producer, start=0):
        """
        Args:
            data_queue(mp.Queue): a queue which contains datapoints.
            nr_producer(int): number of producer processes. This process will
                terminate after receiving this many of :class:`DIE` sentinel.
            start(int): the rank of the first object
        """
        super(OrderedResultGatherProc, self).__init__()
        self.data_queue = data_queue
        self.ordered_container = OrderedContainer(start=start)
        self.result_queue = mp.Queue()
        self.nr_producer = nr_producer

    def run(self):
        nr_end = 0
        try:
            while True:
                task_id, data = self.data_queue.get()
                if task_id == DIE:
                    self.result_queue.put((task_id, data))
                    nr_end += 1
                    if nr_end == self.nr_producer:
                        return
                else:
                    self.ordered_container.put(task_id, data)
                    while self.ordered_container.has_next():
                        self.result_queue.put(self.ordered_container.get())
        except Exception as e:
            import traceback
            traceback.print_exc()
            raise e

    def get(self):
        return self.result_queue.get()