processes_bounded_pool.py

from concurrent.futures import ProcessPoolExecutor
from multiprocessing import BoundedSemaphore
from typing import Optional, Sequence

from concurrent_utils import BoundedPoolExecutor

__all__: Sequence[str] = ("BoundedProcessPoolExecutor",)


class BoundedProcessPoolExecutor(ProcessPoolExecutor, BoundedPoolExecutor):
    """A work pool that uses processes to execute functions asynchronously.

    Importantly, caps the maximum number of executing tasks by using a process-safe semaphore.
    This is useful when one needs to execute many async tasks, but can only have a much smaller
    set of actively executing tasks at once so as to not overload some external system.
    """

    def __init__(self, max_workers: Optional[int] = None) -> None:
        super().__init__(max_workers)
        self.semaphore = BoundedSemaphore(self._max_workers)  # type: ignore