Class ThreadPool

petastorm/workers_pool/thread_pool.py:77–263

class ThreadPool(object):
    def __init__(self, workers_count, results_queue_size=50, profiling_enabled=False, shuffle_rows=False, seed=None):
        """Initializes a thread pool.

        TODO: consider using a standard thread pool
        (e.g. http://elliothallmark.com/2016/12/23/requests-with-concurrent-futures-in-python-2-7/ as an implementation)

        Originally implemented our own pool to match the interface of ProcessPool (could not find a process pool
        implementation that would not use fork)

        :param workers_count: Number of threads
        :param profile: Whether to run a profiler on the threads
        """
        self._seed = seed
        self._shuffle_rows = shuffle_rows
        self._workers = []
        self._ventilator_queues = []
        self.workers_count = workers_count
        self._results_queue_size = results_queue_size
        # Worker threads will watch this event and gracefully shutdown when the event is set
        self._stop_event = Event()
        self._profiling_enabled = profiling_enabled

        # Count of items ventilated by the pool
        self._ventilated_items = 0
        # Count of items ventilated by each worker
        self._ventilated_items_by_worker = [0 for _ in range(self.workers_count)]
        # Count of items processed by each worker
        self._ventilated_items_processed_by_worker = [0 for _ in range(self.workers_count)]
        self._ventilator = None
        self._get_results_worker_id = 0

    def start(self, worker_class, worker_args=None, ventilator=None):
        """Starts worker threads.

        :param worker_class: A class of the worker class. The class will be instantiated in the worker process. The
          class must implement :class:`.WorkerBase` protocol
        :param worker_setup_args: Argument that will be passed to ``args`` property of the instantiated
          :class:`.WorkerBase`
        :return: ``None``
        """
        # Verify stop_event and raise exception if it's already set!
        if self._stop_event.is_set():
            raise RuntimeError('ThreadPool({}) cannot be reused! stop_event set? {}'
                               .format(len(self._workers), self._stop_event.is_set()))

        # Set up a channel for each worker to send work
        self._ventilator_queues = [queue.Queue() for _ in range(self.workers_count)]
        # Set up a channel for each worker to send results
        self._results_queues = [
            queue.Queue(max(5, self._results_queue_size // self.workers_count))
            for _ in range(self.workers_count)
        ]
        self._workers = []
        for worker_id in range(self.workers_count):
            # Create a closure that captures the worker_id for this specific worker
            def make_publish_func(worker_id):
                return lambda data: self._stop_aware_put(worker_id, data)