Map is slow and processes batches one after another

[villmow]

Describe the bug

I have a somewhat unclear bug to me, where I can't figure out what the problem is. The code works as expected on a small subset of my dataset (2000 samples) on my local machine, but when I execute the same code with a larger dataset (1.4 million samples) this problem occurs. Thats why I can't give exact steps to reproduce, I'm sorry.

I process a large dataset in a two step process. I first call map on a dataset I load from disk and create a new dataset from it. This works like expected and map uses all workers I started it with. Then I process the dataset created by the first step, again with map, which is really slow and starting only one or two process at a time. Number of processes is the same for both steps.

pseudo code:

ds = datasets.load_from_disk("path")
new_dataset = ds.map(work, batched=True, ...)  # fast uses all processes
final_dataset = new_dataset.map(work2, batched=True, ...)  # slow starts one process after another

Expected results

Second stage should be as fast as the first stage.

Versions

Paste the output of the following code:

• Datasets: 1.5.0
• Python: 3.8.8 (default, Feb 24 2021, 21:46:12)
• Platform: Linux-5.4.0-60-generic-x86_64-with-glibc2.10

Do you guys have any idea? Thanks a lot!

[albertvillanova]

Hi @villmow, thanks for reporting.

Could you please try with the Datasets version 1.6? We released it yesterday and it fixes some issues about the processing speed. You can see the fix implemented by @lhoestq here: #2122.

Once you update Datasets, please confirm if the problem persists.

[villmow]

Hi @albertvillanova, thanks for the reply. I just tried the new version and the problem still persists.

Do I need to rebuild the saved dataset (which I load from disk) with the 1.6.0 version of datasets? My script loads this dataset and creates new datasets from it. I tried it without rebuilding.

See this short video of what happens. It does not create all processes at the same time:

problem.mp4

[lhoestq]

There can be a bit of delay between the creations of the processes but this delay should be the same for both your map calls. We should look into this.
Also if you hav some code that reproduces this issue on google colab that'd be really useful !

Regarding the speed differences:
This looks like a similar issue as #1992 who is experiencing the same speed differences between processes.
This is a known bug that we are investigating. As of now I've never managed to reproduce it on my machine so it's pretty hard for me to find where this issue comes from.

[villmow]

Upgrade to 1.6.1 solved my problem somehow. I did not change any of my code, but now it starts all processes around the same time.

[lhoestq]

Nice ! I'm glad this works now.
Closing for now, but feel free to re-open if you experience this issue again.