`Dataset.save_to_disk` hangs when using num_proc > 1

[JohannesAck]

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 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

[GeekOrangeLuYao]

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.

[Macchiatobin]

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