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

Code before:
    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():
                # we get the result in case there's an error to raise
                [async_result.get(timeout=0.05) for async_result in async_results]


Review comment (bug):
Could we do this automatically outside of `iflatmap_unordered` ? There is a thread is automatically repopulates the processes that crashed. Maybe you can just override the repopulation to raise an error if it's called more than the size of the pool times.

Code after:
    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]