Repository: huggingface/datasets
File: src/datasets/utils/py_utils.py (Python)

Code before:
    """
    A multiprocess.pool.Pool implementation that keeps track of child process' PIDs,
    and can detect if a child process has been restarted.
    """

    def __init__(self, *args, **kwargs) -> None:
        super().__init__(*args, **kwargs)
        self._last_pool_pids = self._get_current_pool_pids()
        self._has_restarted_subprocess = False

    def _get_current_pool_pids(self) -> Set[int]:
        return {f.pid for f in self._pool}

    def has_restarted_subprocess(self) -> bool:
        if self._has_restarted_subprocess:
            # If the pool ever restarted a subprocess,
            # we don't check the PIDs again.
            return True
        current_pids = self._get_current_pool_pids()
        self._has_restarted_subprocess = current_pids != self._last_pool_pids
        self._last_pool_pids = current_pids
        return self._has_restarted_subprocess


def iflatmap_unordered(
    pool: ProcessPool,
    func: Callable[..., Iterable[Y]],
    *,
    kwargs_iterable: Iterable[dict],
) -> Iterable[Y]:
    manager_cls = Manager if isinstance(pool, multiprocessing.pool.Pool) else multiprocess.Manager
    with manager_cls() as manager:
        queue = manager.Queue()
        async_results = [
            pool.apply_async(_write_generator_to_queue, (queue, func, kwargs)) for kwargs in kwargs_iterable
        ]
        try:
            while True:
                try:
                    yield queue.get(timeout=0.05)
                except Empty:
                    if all(async_result.ready() for async_result in async_results) and queue.empty():
                        break
                if pool.has_restarted_subprocess():
                    # One of the subprocesses has died. We should not wait forever.
                    raise RuntimeError(
                        "One of the subprocesses has abruptly died during map operation."
                        "To debug the error, disable multiprocessing."
                    )
        finally:
            if not pool.has_restarted_subprocess():

Review comment (suggestion):
I think we still need to keep `Union[multiprocessing.pool.Pool, multiprocess.pool.Pool]` support in `iflatmap_unordered`

Code after:
    return {f.pid for f in pool._pool}


def iflatmap_unordered(
    pool: Union[multiprocessing.pool.Pool, multiprocess.pool.Pool],
    func: Callable[..., Iterable[Y]],
    *,
    kwargs_iterable: Iterable[dict],
) -> Iterable[Y]:
    initial_pool_pid = _get_pool_pid(pool)
    pool_changed = False
    manager_cls = Manager if isinstance(pool, multiprocessing.pool.Pool) else multiprocess.Manager
    with manager_cls() as manager:
        queue = manager.Queue()
        async_results = [
            pool.apply_async(_write_generator_to_queue, (queue, func, kwargs)) for kwargs in kwargs_iterable
        ]
        try:
            while True:
                try:
                    yield queue.get(timeout=0.05)
                except Empty:
                    if all(async_result.ready() for async_result in async_results) and queue.empty():
                        break
                if _get_pool_pid(pool) != initial_pool_pid:
                    pool_changed = True
                    # One of the subprocesses has died. We should not wait forever.
                    raise RuntimeError(
                        "One of the subprocesses has abruptly died during map operation."
                        "To debug the error, disable multiprocessing."
                    )
        finally:
            if not pool_changed:
                # we get the result in case there's an error to raise
                [async_result.get(timeout=0.05) for async_result in async_results]