`Dataset.save_to_disk` hangs when using num_proc > 1
Opened by JohannesAck (open)

### Describe the bug

Hi, I'm encountered a small issue when saving datasets that led to the saving taking up to multiple hours.
Specifically, [`Dataset.save_to_disk`](https://huggingface.co/docs/datasets/main/en/package_reference/main_classes#datasets.Dataset.save_to_disk) is a lot slower when using `num_proc>1` than when using `num_proc=1`

The documentation mentions that "Multiprocessing is disabled by default.", but there is no explanation on how to enable it.

### Steps to reproduce the bug

```
import numpy as np
from datasets import Dataset

n_samples = int(4e6)
n_tokens_sample = 100
data_dict = {
    'tokens' : np.random.randint(0, 100, (n_samples, n_tokens_sample)),
}

dataset = Dataset.from_dict(data_dict)
dataset.save_to_disk('test_dataset', num_proc=1)
dataset.save_to_disk('test_dataset', num_proc=4)
dataset.save_to_disk('test_dataset', num_proc=8)
```

This results in:
```
>>> dataset.save_to_disk('test_dataset', num_proc=1)
Saving the dataset (7/7 shards): 100%|██████████████| 4000000/4000000 [00:17<00:00, 228075.15 examples/s]
>>> dataset.save_to_disk('test_dataset', num_proc=4)
Saving the dataset (7/7 shards): 100%|██████████████| 4000000/4000000 [01:49<00:00, 36583.75 examples/s]
>>> dataset.save_to_disk('test_dataset', num_proc=8)
Saving the dataset (8/8 shards): 100%|██████████████| 4000000/4000000 [02:11<00:00, 30518.43 examples/s]
```

With larger datasets it can take hours, but I didn't benchmark that for this bug report.

### Expected behavior

I would expect using `num_proc>1` to be faster instead of slower than `num_proc=1`. 

### Environment info

- `datasets` version: 3.1.0
- Platform: Linux-5.15.153.1-microsoft-standard-WSL2-x86_64-with-glibc2.35
- Python version: 3.10.12
- `huggingface_hub` version: 0.26.2
- PyArrow version: 18.0.0
- Pandas version: 2.2.3
- `fsspec` version: 2024.6.1

Comments:
I've met the same situations.

Here's my logs:
num_proc = 64, I stop it early as it cost **too** much time.
```
Saving the dataset (1540/4775 shards):  32%|███▏      | 47752224/147853764 [15:32:54<132:28:34, 209.89 examples/s]
Saving the dataset (1540/4775 shards):  32%|███▏      | 47754224/147853764 [15:32:54<78:14:37, 355.37 examples/s] 
Saving the dataset (1540/4775 shards):  32%|███▏      | 47755224/147853764 [15:32:59<93:43:45, 296.65 examples/s]
```

num_proc = 1(Not set num_proc parameter)
```
Saving the dataset (1753/4775 shards):  37%|███▋      | 54301556/147853764 [24:46<38:33, 40440.93 examples/s]
Saving the dataset (1753/4775 shards):  37%|███▋      | 54306556/147853764 [24:46<39:34, 39392.01 examples/s]
Saving the dataset (1753/4775 shards):  37%|███▋      | 54311520/147853764 [24:46<38:56, 40030.53 examples/s]
```

I check the conditions of CPUs and Memory I/O, I found that disk I/O was blocked, but CPU and memory usage was high. There should be some bugs in the code.



Any new process on this issue? I'm encountering the same issue.

Not getting this issue. 

My output;

`Saving the dataset (0/87 shards):   7%|▎    | 588000/8557560 [01:29<21:02, 6314.04 examples/s]`

At setting `num_proc=64`, and 

`Saving the dataset (0/87 shards):   0%|    | 28000/8557560 [03:20<16:59:06, 139.49 examples/s]` 

At num_proc=1 (pass nothing)

My `pyproject.toml`; 

```
[project]
name = "test"
version = "0.1.0"
requires-python = "==3.13.0"
dependencies = [
    "absl-py==2.1.0",
    "accelerate==1.7.0",
    "aiohappyeyeballs==2.6.1",
    "aiohttp==3.12.11",
    "aiosignal==1.3.2",
    "annotated-types==0.7.0",
    "appdirs==1.4.4",
    "argcomplete>=1.8.1",
    "astunparse==1.6.3",
    "async-timeout==5.0.1",
    "attrs==21.2.0",
    "automat==20.2.0",
    "babel==2.8.0",
    "backcall==0.2.0",
    "bcrypt==3.2.0",
    "beautifulsoup4==4.10.0",
    "beniget==0.4.2",
    "bleach==4.1.0",
    "blinker==1.4",
    "blis==1.3.0",
    "bottle==0.12.19",
    "brotli==1.0.9",
    "catalogue==2.0.10",
    "certifi==2020.6.20",
    "cffi==1.15.0",
    "chardet==4.0.0",
    "charset-normalizer==3.4.2",
    "click==8.0.3",
    "cloudpathlib==0.21.1",
    "colorama>=0.4.4",
    "commonmark==0.9.1",
    "confection==0.1.5",
    "configobj==5.0.6",
    "constantly==15.1.0",
    "cryptography==3.4.8",
    "ctop==1.0.0",
    "cycler==0.11.0",
    "cymem==2.0.11",
    "datasets==3.6.0",
    "dbus-python==1.2.18",
    "decorator==4.4.2",
    "defusedxml==0.7.1",
    "dill==0.3.8",
    "distlib==0.3.4",
    "distro==1.7.0",
    "einops==0.8.1",
    "en-core-web-sm",
    "entrypoints==0.4",
    "evaluate==0.4.3",
    "filelock==3.6.0",
    "flake8==4.0.1",
    "flatbuffers==25.2.10",
    "fonttools==4.29.1",
    "frozenlist==1.7.0",
    "fs==2.4.12",
    "fsspec==2024.3.1",
    "future==0.18.2",
    "gast==0.6.0",
    "gitdb==4.0.12",
    "gitpython==3.1.44",
    "glances==3.2.4.2",
    "google-pasta==0.2.0",
    "grpcio==1.73.0",
    "h5py==3.14.0",
    "hf-xet==1.1.3",
    "html5lib==1.1",
    "httplib2==0.20.2",
    "huggingface-hub==0.32.4",
    "hyperlink==21.0.0",
    "icdiff==2.0.4",
    "idna==3.3",
    "importlib-metadata==4.6.4",
    "incremental==21.3.0",
    "influxdb==5.3.1",
    "ipykernel==6.7.0",
    "ipython==7.31.1",
    "ipython-genutils==0.2.0",
    "jax>=0.5.1",
    "jax-cuda12-pjrt>=0.5.1",
    "jax-cuda12-plugin>=0.5.1",
    "jaxlib>=0.5.1",
    "jedi==0.18.0",
    "jeepney==0.7.1",
    "jinja2==3.0.3",
    "joblib>=0.17.0",
    "jsonpatch==1.32",
    "jsonpointer==2.0",
    "jsonschema==3.2.0",
    "jupyter-client==7.1.2",
    "jupyter-core==4.9.1",
    "kaptan>=0.5.12",
    "keras==3.6.0",
    "keyring==23.5.0",
    "kiwisolver==1.3.2",
    "langcodes==3.5.0",
    "language-data==1.3.0",
    "launchpadlib==1.10.16",
    "lazr-restfulclient==0.14.4",
    "lazr-uri==1.0.6",
    "libclang==18.1.1",
    "libtmux==0.10.1",
    "livereload==2.6.3",
    "lxml==6.0.0",
    "lz4==3.1.3",
    "marisa-trie==1.2.1",
    "markdown==3.3.6",
    "markupsafe==2.0.1",
    "matplotlib==3.5.1",
    "matplotlib-inline==0.1.3",
    "mccabe==0.6.1",
    "mkdocs==1.1.2",
    "ml-dtypes==0.5.1",
    "more-itertools==8.10.0",
    "mpmath>=0.0.0",
    "msgpack==1.0.3",
    "multidict==6.4.4",
    "multiprocess==0.70.16",
    "murmurhash==1.0.13",
    "namex==0.0.8",
    "nest-asyncio==1.5.4",
    "netifaces==0.11.0",
    "networkx==2.4",
    "nltk==3.9.1",
    "numpy>=1.23.3",
    "nvidia-ml-py==12.555.43",
    "oauthlib==3.2.0",
    "olefile==0.46",
    "opt-einsum==3.3.0",
    "optree==0.13.1",
    "packaging==21.3",
    "pandas==2.3.0",
    "parso==0.8.1",
    "pexpect==4.8.0",
    "pickleshare==0.7.5",
    "pillow>=9.0.1",
    "pipx==1.0.0",
    "platformdirs==2.5.1",
    "ply==3.11",
    "portalocker==3.1.1",
    "preshed==3.0.10",
    "prompt-toolkit==3.0.28",
    "propcache==0.3.2",
    "protobuf==4.21.12",
    "psutil==5.9.0",
    "ptyprocess==0.7.0",
    "py==1.10.0",
    "pyarrow==20.0.0",
    "pyasn1==0.4.8",
    "pyasn1-modules==0.2.1",
    "pycodestyle==2.8.0",
    "pycparser==2.21",
    "pycryptodomex==3.11.0",
    "pydantic==2.11.5",
    "pydantic-core==2.33.2",
    "pyflakes==2.4.0",
    "pygments==2.11.2",
    "pyhamcrest==2.0.2",
    "pyinotify==0.9.6",
    "pyjwt==2.3.0",
    "pyopenssl==21.0.0",
    "pyparsing==2.4.7",
    "pyrsistent==0.18.1",
    "pyserial==3.5",
    "pysmi==0.3.2",
    "pysnmp==4.4.12",
    "pystache==0.6.0",
    "python-dateutil>=2.8.1",
    "python-magic==0.4.24",
    "pythran>=0.10.0",
    "pytz==2022.1",
    "pyyaml==6.0.2",
    "regex==2024.11.6",
    "requests==2.32.4",
    "rich==11.2.0",
    "rouge-score==0.1.2",
    "sacrebleu==2.5.1",
    "safetensors==0.5.3",
    "scikit-learn==1.7.0",
    "scipy>=1.8.0",
    "secretstorage==3.3.1",
    "sentence-transformers==4.1.0",
    "sentry-sdk==2.29.1",
    "service-identity==18.1.0",
    "setproctitle==1.3.6",
    "shellingham==1.5.4",
    "six==1.16.0",
    "smart-open==7.1.0",
    "smmap==5.0.2",
    "soupsieve==2.3.1",
    "spacy==3.8.7",
    "spacy-legacy==3.0.12",
    "spacy-loggers==1.0.5",
    "srsly==2.5.1",
    "ssh-import-id==5.11",
    "sympy>=1.12",
    "tabulate==0.9.0",
    "tensorboard==2.19.0",
    "tensorboard-data-server==0.7.2",
    "termcolor==1.1.0",
    "thinc==8.3.6",
    "threadpoolctl==3.1.0",
    "tmuxp==1.9.2",
    "tokenizers==0.21.1",
    "torch==2.6.0",
    "torchvision==0.21.0",
    "tornado==6.1",
    "tqdm==4.67.1",
    "traitlets==5.1.1",
    "transformers==4.52.4",
    "triton==3.2.0",
    "twisted==22.1.0",
    "typer==0.16.0",
    "typing-extensions==4.14.0",
    "typing-inspection==0.4.1",
    "ufolib2==0.13.1",
    "urllib3==2.4.0",
    "userpath==1.8.0",
    "virtualenv==20.13.0",
    "wadllib==1.3.6",
    "wandb==0.20.1",
    "wasabi==1.1.3",
    "wcwidth==0.2.5",
    "weasel==0.4.1",
    "webencodings==0.5.1",
    "websocket-client==1.2.3",
    "werkzeug==2.0.2",
    "xxhash==3.5.0",
    "yarl==1.20.1",
    "zipp==1.0.0",
    "zope-interface==5.4.0",
]

[tool.uv.sources]
en-core-web-sm = { url = "https://github.com/explosion/spacy-models/releases/download/en_core_web_sm-3.8.0/en_core_web_sm-3.8.0-py3-none-any.whl" }
```