Fix parallel downloads for datasets without scripts (#6551)

* fix multiprocessing download

* comment

Author: lhoestq

src/datasets/download/download_manager.py

@@ -30,7 +30,14 @@
 from .. import config
 from ..utils import tqdm as hf_tqdm
 from ..utils.deprecation_utils import DeprecatedEnum, deprecated
-from ..utils.file_utils import cached_path, get_from_cache, hash_url_to_filename, is_relative_path, url_or_path_join
+from ..utils.file_utils import (
+    cached_path,
+    get_from_cache,
+    hash_url_to_filename,
+    is_relative_path,
+    stack_multiprocessing_download_progress_bars,
+    url_or_path_join,
+)
 from ..utils.info_utils import get_size_checksum_dict
 from ..utils.logging import get_logger
 from ..utils.py_utils import NestedDataStructure, map_nested, size_str
@@ -423,13 +430,14 @@ def download(self, url_or_urls):
         download_func = partial(self._download, download_config=download_config)
 
         start_time = datetime.now()
-        downloaded_path_or_paths = map_nested(
-            download_func,
-            url_or_urls,
-            map_tuple=True,
-            num_proc=download_config.num_proc,
-            desc="Downloading data files",
-        )
+        with stack_multiprocessing_download_progress_bars():
+            downloaded_path_or_paths = map_nested(
+                download_func,
+                url_or_urls,
+                map_tuple=True,
+                num_proc=download_config.num_proc,
+                desc="Downloading data files",
+            )
         duration = datetime.now() - start_time
         logger.info(f"Downloading took {duration.total_seconds() // 60} min")
         url_or_urls = NestedDataStructure(url_or_urls)

src/datasets/utils/file_utils.py

@@ -7,6 +7,7 @@
 import copy
 import io
 import json
+import multiprocessing
 import os
 import posixpath
 import re
@@ -19,6 +20,7 @@
 from functools import partial
 from pathlib import Path
 from typing import Optional, TypeVar, Union
+from unittest.mock import patch
 from urllib.parse import urljoin, urlparse
 
 import fsspec
@@ -319,6 +321,12 @@ def fsspec_head(url, storage_options=None):
     return fs.info(paths[0])
 
 
+def stack_multiprocessing_download_progress_bars():
+    # Stack downloads progress bars automatically using HF_DATASETS_STACK_MULTIPROCESSING_DOWNLOAD_PROGRESS_BARS=1
+    # We use environment variables since the download may happen in a subprocess
+    return patch.dict(os.environ, {"HF_DATASETS_STACK_MULTIPROCESSING_DOWNLOAD_PROGRESS_BARS": "1"})
+
+
 class TqdmCallback(fsspec.callbacks.TqdmCallback):
     def __init__(self, tqdm_kwargs=None, *args, **kwargs):
         super().__init__(tqdm_kwargs, *args, **kwargs)
@@ -335,6 +343,10 @@ def fsspec_get(url, temp_file, storage_options=None, desc=None):
             "desc": desc or "Downloading",
             "unit": "B",
             "unit_scale": True,
+            "position": multiprocessing.current_process()._identity[-1]  # contains the ranks of subprocesses
+            if os.environ.get("HF_DATASETS_STACK_MULTIPROCESSING_DOWNLOAD_PROGRESS_BARS") == "1"
+            and multiprocessing.current_process()._identity
+            else None,
         }
     )
     fs.get_file(paths[0], temp_file.name, callback=callback)
@@ -389,6 +401,10 @@ def http_get(
         total=total,
         initial=resume_size,
         desc=desc or "Downloading",
+        position=multiprocessing.current_process()._identity[-1]  # contains the ranks of subprocesses
+        if os.environ.get("HF_DATASETS_STACK_MULTIPROCESSING_DOWNLOAD_PROGRESS_BARS") == "1"
+        and multiprocessing.current_process()._identity
+        else None,
     ) as progress:
         for chunk in response.iter_content(chunk_size=1024):
             progress.update(len(chunk))

src/datasets/utils/py_utils.py

@@ -462,7 +462,18 @@ def map_nested(
 
     if num_proc is None:
         num_proc = 1
-    if num_proc != -1 and num_proc <= 1 or len(iterable) < parallel_min_length:
+    if any(isinstance(v, types) and len(v) > len(iterable) for v in iterable):
+        mapped = [
+            map_nested(
+                function=function,
+                data_struct=obj,
+                num_proc=num_proc,
+                parallel_min_length=parallel_min_length,
+                types=types,
+            )
+            for obj in iterable
+        ]
+    elif num_proc != -1 and num_proc <= 1 or len(iterable) < parallel_min_length:
         mapped = [
             _single_map_nested((function, obj, types, None, True, None))
             for obj in hf_tqdm(iterable, disable=disable_tqdm, desc=desc)