Creating a process Pool subclass to track PID changes

Author: pappacena

src/datasets/arrow_dataset.py

@@ -58,7 +58,6 @@
 import pyarrow as pa
 import pyarrow.compute as pc
 from huggingface_hub import HfApi, HfFolder
-from multiprocess import Pool
 from requests import HTTPError
 
 from . import config
@@ -113,7 +112,7 @@
 from .utils.hub import hf_hub_url
 from .utils.info_utils import is_small_dataset
 from .utils.metadata import DatasetMetadata
-from .utils.py_utils import Literal, asdict, convert_file_size_to_int, iflatmap_unordered, unique_values
+from .utils.py_utils import Literal, ProcessPool, asdict, convert_file_size_to_int, iflatmap_unordered, unique_values
 from .utils.stratify import stratified_shuffle_split_generate_indices
 from .utils.tf_utils import dataset_to_tf, minimal_tf_collate_fn, multiprocess_dataset_to_tf
 from .utils.typing import ListLike, PathLike
@@ -1505,7 +1504,7 @@ def save_to_disk(
         shard_lengths = [None] * num_shards
         shard_sizes = [None] * num_shards
         if num_proc > 1:
-            with Pool(num_proc) as pool:
+            with ProcessPool(num_proc) as pool:
                 with pbar:
                     for job_id, done, content in iflatmap_unordered(
                         pool, Dataset._save_to_disk_single, kwargs_iterable=kwargs_per_job
@@ -3167,7 +3166,7 @@ def format_new_fingerprint(new_fingerprint: str, rank: int) -> str:
                     logger.info(
                         f"Reprocessing {len(kwargs_per_job)}/{num_shards} shards because some of them were missing from the cache."
                     )
-                with Pool(len(kwargs_per_job)) as pool:
+                with ProcessPool(len(kwargs_per_job)) as pool:
                     os.environ = prev_env
                     logger.info(f"Spawning {num_proc} processes")
                     with logging.tqdm(

src/datasets/builder.py

@@ -33,7 +33,6 @@
 
 import fsspec
 import pyarrow as pa
-from multiprocess import Pool
 from tqdm.contrib.concurrent import thread_map
 
 from . import config, utils
@@ -69,6 +68,7 @@
 from .utils.filelock import FileLock
 from .utils.info_utils import VerificationMode, get_size_checksum_dict, verify_checksums, verify_splits
 from .utils.py_utils import (
+    ProcessPool,
     classproperty,
     convert_file_size_to_int,
     has_sufficient_disk_space,
@@ -1532,7 +1532,7 @@ def _prepare_split(
             shards_per_job = [None] * num_jobs
             shard_lengths_per_job = [None] * num_jobs
 
-            with Pool(num_proc) as pool:
+            with ProcessPool(num_proc) as pool:
                 with pbar:
                     for job_id, done, content in iflatmap_unordered(
                         pool, self._prepare_split_single, kwargs_iterable=kwargs_per_job
@@ -1791,7 +1791,7 @@ def _prepare_split(
             shards_per_job = [None] * num_jobs
             shard_lengths_per_job = [None] * num_jobs
 
-            with Pool(num_proc) as pool:
+            with ProcessPool(num_proc) as pool:
                 with pbar:
                     for job_id, done, content in iflatmap_unordered(
                         pool, self._prepare_split_single, kwargs_iterable=kwargs_per_job

src/datasets/utils/py_utils.py

@@ -1330,39 +1330,57 @@ def _write_generator_to_queue(queue: queue.Queue, func: Callable[..., Iterable[Y
     return i
 
 
-def _get_pool_pid(pool: Union[multiprocessing.pool.Pool, multiprocess.pool.Pool]) -> Set[int]:
-    return {f.pid for f in pool._pool}
+class ProcessPool(multiprocess.pool.Pool):
+    """
+    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: Union[multiprocessing.pool.Pool, multiprocess.pool.Pool],
+    pool: ProcessPool,
     func: Callable[..., Iterable[Y]],
     *,
     kwargs_iterable: Iterable[dict],
 ) -> Iterable[Y]:
-    initial_pool_pid = _get_pool_pid(pool)
     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
         ]
-        subproc_killed = False
         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:
-                    subproc_killed = True
+                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 subproc_killed:
+            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]

tests/test_arrow_dataset.py

@@ -1343,9 +1343,9 @@ def test_map_caching(self, in_memory):
             with self._caplog.at_level(WARNING):
                 with self._create_dummy_dataset(in_memory, tmp_dir) as dset:
                     with patch(
-                        "datasets.arrow_dataset.Pool",
+                        "datasets.arrow_dataset.ProcessPool",
                         new_callable=PickableMagicMock,
-                        side_effect=datasets.arrow_dataset.Pool,
+                        side_effect=datasets.arrow_dataset.ProcessPool,
                     ) as mock_pool:
                         with dset.map(lambda x: {"foo": "bar"}, num_proc=2) as dset_test1:
                             dset_test1_data_files = list(dset_test1.cache_files)