Class ProcessPool

petastorm/workers_pool/process_pool.py:114–312

class ProcessPool(object):
    def __init__(self, workers_count, serializer=None, zmq_copy_buffers=True):
        """Initializes a ProcessPool.

        This pool is different from standard Python pool implementations by the fact that the workers are spawned
        without using fork. Some issues with using jvm based HDFS driver were observed when the process was forked
        (could not access HDFS from the forked worker if the driver was already used in the parent process).

        :param workers_count: Number of processes to be spawned
        :param serializer: An object that would be used for data payload serialization when sending data from a worker
          process to the main process. ``PickleSerializer`` is used by default. May use
          :class:`petastorm.reader_impl.ArrowTableSerializer` (should be used together with
          :class:`petastorm.reader.ArrowReader`)
        :param zmq_copy_buffers: When set to False, we will use a zero-memory-copy feature of recv_multipart.
          A downside of using this zero memory copy feature is that it does not play nice with Python GC and cases
          were observed when it resulted in wild memory footprint swings. Having the buffers copied is typically a
          safer alternative.
        """
        self._workers = []
        self._ventilator_send = None
        self._control_sender = None
        self.workers_count = workers_count
        self._results_receiver_poller = None
        self._results_receiver = None

        self._ventilated_items = 0
        self._ventilated_items_processed = 0
        self._ventilator = None
        self._serializer = serializer or PickleSerializer()
        self._zmq_copy_buffers = zmq_copy_buffers

    def _create_local_socket_on_random_port(self, context, socket_type):
        """Creates a zmq socket on a random port.

        :param context: zmq context
        :param socket_type: zmq socket type
        :return: A tuple: ``(zmq_socket, endpoint_address)``
        """
        LOCALHOST = 'tcp://127.0.0.1'
        socket = context.socket(socket_type)

        # There are race conditions where the socket can close when messages are still trying to be sent by zmq.
        # This can end up causing zmq to block indefinitely when sending objects or shutting down. Having the socket
        # linger on close helps prevent this.
        socket.linger = _SOCKET_LINGER_MS

        port = socket.bind_to_random_port(LOCALHOST)
        return socket, '{}:{}'.format(LOCALHOST, port)

    def start(self, worker_class, worker_setup_args=None, ventilator=None):
        """Starts worker processes.

        Will block until all processes to subscribe to the worker queue (the messages are distributed by zmq on write
        so if only one, out of many, workers is up at the time of 'ventilation', the initial load won't be balanced
        between workers. If can not start the workers in timely fashion, will raise an exception.

        :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.