Support DataLoader with num_workers > 0 in streaming mode
Opened by lhoestq (closed)

### Issue

It's currently not possible to properly stream a dataset using multiple `torch.utils.data.DataLoader` workers:

- the `TorchIterableDataset` can't be pickled and passed to the subprocesses: https://github.com/huggingface/datasets/issues/3950
- streaming extension is failing: https://github.com/huggingface/datasets/issues/3951
- `fsspec` doesn't work out of the box in subprocesses

### Solution in this PR

I fixed these to enable passing an `IterableDataset` to a  `torch.utils.data.DataLoader` with `num_workers > 0`.

I also had to shard the `IterableDataset` to give each worker a shard, otherwise data would be duplicated. This is implemented in `TorchIterableDataset.__iter__` and uses the new `IterableDataset._iter_shard(shard_idx)` method

I also had to do a few changes the patching that enable streaming in dataset scripts:
- the patches are now always applied - not just for streaming mode. They're applied when a builder is instantiated
- I improved it to also check for renamed modules or attributes (ex: pandas vs pd)
- I grouped all the patches of pathlib.Path into a class `xPath`, so that `Path` outside of dataset scripts stay unchanged - otherwise I didn't change the content of the extended Path methods for streaming
- I fixed a bug with the `pd.read_csv` patch, opening the file in "rb" mode was missing and causing some datasets to not work in streaming mode, and compression inference was missing

### A few details regarding `fsspec` in multiprocessing

From https://github.com/fsspec/filesystem_spec/pull/963#issuecomment-1131709948 :
> Non-async instances might be safe in the forked child, if they hold no open files/sockets etc.; I'm not sure any implementations pass this test!
> If any async instance has been created, the newly forked processes must:
> 1. discard references to locks, threads and event loops and make new ones
> 2. not use any async fsspec instances from the parent process
> 3. clear all class instance caches

Therefore in a DataLoader's worker, I clear the reference to the loop and thread (1). We should be fine for 2 and 3 already since we don't use fsspec class instances from the parent process.

Fix https://github.com/huggingface/datasets/issues/3950
Fix https://github.com/huggingface/datasets/issues/3951

TODO:
- [x] fix tests

Comments:
_The documentation is not available anymore as the PR was closed or merged._

Alright this is finally ready for review ! It's quite long I'm sorry, but it's not easy to disentangle everything ^^'

The main additions are in
- src/datasets/formatting/dataset_wrappers/torch_iterable_dataset.py
- src/datasets/iterable_dataset.py
- src/datasets/utils/patching.py

Added some comments and an error when lists have different lengths for sharding :)

Let's resolve the merge conflict and the CI error (if it's related to the changes), and I can review the PR again.

Feel free to review again :) The CI fail is unrelated to this PR and will be fixed by https://github.com/huggingface/datasets/pull/4472 (the hub now returns 401 instead of 404 for unauthenticated requests to non-existing repos)

CI failures are unrelated to this PR - merging :)

(CI fails are a mix of pip install fails and Hub fails)

@lhoestq you're our hero :)