Attempt to fix multiprocessing hang by closing and joining the pool before termination

[dakinggg]

#6393 has plagued me on and off for a very long time. I have had various workarounds (one time combining two filter calls into one filter call removed the issue, another time making rank 0 go first resolved a cache race condition, one time i think upgrading the version of something resolved it). I don't know hf datasets well enough to fully understand the root cause, but I think this PR fixes it.

Evidence: I have an LLM Foundry training yaml/script (datasets version 3.2.0) that results in a hang ~1/10 times (for a baseline for this testing, it was 2/36 runs that hung). I also reran with the latest datasets version (3.3.1) and got 4/36 hung. Installing datasets from this PR, I was able to successful run the script 144 times without a hang occurring. Assuming the base probability is 1/10, this should be more than enough times to have confidence it works.

After adding some logging, I could see that the code hung during the exit of the mp pool context manager, after all shards had been processed, and the tqdm context manager had exited.

My best explanation: When multiprocessing pool exit is called, it calls pool.terminate, which forcefully exits all the processes (and calls code related to this that I haven't looked at closely). I'm guessing this forceful termination has a bad interaction with some multithreading/multiprocessing that hf datasets does. If we instead call pool.close and pool.join before the pool.terminate happens, perhaps whatever that bad interaction is is able to complete gracefully, and then terminate call proceeds without issue.

If this PR seems good to you, I'd be very appreciative if you were able to do a patch release including it. Thank you!

@lhoestq

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[lhoestq]

Hi Daniel, sure I'll do a patch release soon

[lorabit110]

Thanks for the fix! We have been affected by this a lot when we try to use LLM Foundry with larger multimodal ICL datasets.

[dakinggg]

@lorabit110 are you able to test it out for your case as well? Would be great to get a second validation that it actually fixes the issue. Thanks!