Support DataLoader with num_workers > 0 in streaming mode #4375

lhoestq · 2022-05-19T15:00:31Z

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: Streaming Datasets don't work with Transformers Trainer when dataloader_num_workers>1 #3950
streaming extension is failing: Forked streaming datasets try to open data urls rather than use network #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 fsspec/filesystem_spec#963 (comment) :

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:

discard references to locks, threads and event loops and make new ones

not use any async fsspec instances from the parent process

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 #3950
Fix #3951

TODO:

fix tests

- make it picklable - paralellize over the shards when num_workers is passed

HuggingFaceDocBuilderDev · 2022-05-19T15:14:01Z

The documentation is not available anymore as the PR was closed or merged.

lhoestq · 2022-06-07T13:18:04Z

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

lhoestq · 2022-06-07T14:33:00Z

src/datasets/download/streaming_download_manager.py



-def xpathrglob(path, pattern, **kwargs):
-    """Rglob function for argument of type :obj:`~pathlib.Path` that supports both local paths end remote URLs.
+class xPath(type(Path())):


many changes in this file are just about moving functions inside this class.

For example I moved xpathrglob to xPath.rglob

lhoestq · 2022-06-07T15:13:35Z

src/datasets/formatting/dataset_wrappers/torch_iterable_dataset.py

+            if worker_info.id == 0 and self.n_shards < worker_info.num_workers:
+                logger.warning(
+                    f"Too many dataloader workers: {worker_info.num_workers} (max is dataset.n_shards={self.n_shards}). "
+                    f"Stopping dataloader workers [{self.n_shards}...{worker_info.num_workers -1}]."
+                )
+                logger.warning(
+                    f"To parallelize data loading, we give each process some shards (or data sources) to process. "
+                    f"Therefore it's unnecessary to have a number of workers greater than dataset.n_shards={self.n_shards}."
+                    f"To enable more parallelism, please split the dataset in more files than {self.n_shards}."
+                )
+            # split workload
+            shards_indices = list(range(worker_info.id, self.n_shards, worker_info.num_workers))
+            if shards_indices:
+                logger.debug(
+                    f"dataloader worker#{worker_info.id}, ': Starting to iterate over {len(shards_indices)}/{self.n_shards} shards."
+                )
+                for shard_idx in shards_indices:
+                    for key, example in self._iter_shard(shard_idx):
+                        yield self._apply_feature_types(example)


this is where we shard the iterable dataset when it's passed to a DataLoader worker

lhoestq · 2022-06-07T15:15:51Z

src/datasets/iterable_dataset.py

+            ex_iterable = self._ex_iterable.shuffle_data_sources(self._effective_generator())
+        else:
+            ex_iterable = self._ex_iterable
+        yield from ex_iterable.shard_data_sources(shard_idx)


This is what is called when iterating in a DataLoader worker. The idea is to iterate only on one shard out of the self.n_shards available

lhoestq · 2022-06-07T15:18:47Z

src/datasets/streaming.py

-    # xml.etree.ElementTree
-    for submodule in ["ElementTree", "ET"]:
-        patch_submodule(module, f"{submodule}.parse", wrap_auth(xet_parse)).start()
+    patch_submodule(module, "pathlib.Path", xPath).start()


Now we can just pass the source object to patch in the module, and it will patch it even if the attribute alone is imported, or even if a parent module has been imported and renamed (see test_patching.py for a list of all supported cases - I probably have to add a docstring to patch_submodule as well)

mariosasko

This is 🔥. Just two comments:

(Also thanks for the the review instructions/code clarifications)

src/datasets/iterable_dataset.py

src/datasets/utils/patching.py

lhoestq · 2022-06-09T10:03:57Z

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

mariosasko · 2022-06-09T10:44:26Z

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

lhoestq · 2022-06-10T12:30:15Z

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

mariosasko

All looks good now! Thanks!

src/datasets/iterable_dataset.py

Co-authored-by: Mario Šaško <mariosasko777@gmail.com>

lhoestq · 2022-06-10T20:47:18Z

CI failures are unrelated to this PR - merging :)

(CI fails are a mix of pip install fails and Hub fails)

justheuristic · 2022-07-04T16:05:13Z

@lhoestq you're our hero :)

lhoestq added 4 commits April 4, 2022 19:11

make TorchIterableDataset work in parallel

58289fc

- make it picklable - paralellize over the shards when num_workers is passed

start writing some tests

c26a9f1

Merge branch 'master' into parallel-torch-iterable-dataset

4c3ce96

fix streaming extension and fsspec issues in subprocesses

8c60fa3

lhoestq added 14 commits May 19, 2022 17:59

fix some tests

6dc859e

fix more tests

7056f1a

Merge branch 'master' into parallel-torch-iterable-dataset

edef69b

fix import

c0a0492

fix and add tests

7043816

fix patch (handle successive patches and builtins)

a9ea955

revert unnecessary change to enriched_web_blg

07d4c0e

style

af5de1a

use open locally to fix win permission errors

b84ae0e

keep file opened in read_csv

1746712

Merge branch 'master' into parallel-torch-iterable-dataset

bc837ce

fix compression for read_csv

fe269bf

consistency of read_csv: don't infer compression for file-like objects

482c4fb

stringify Path objects

54e9f39

lhoestq marked this pull request as ready for review June 7, 2022 11:56

lhoestq requested a review from mariosasko June 7, 2022 11:56

lhoestq commented Jun 7, 2022

View reviewed changes

mariosasko reviewed Jun 8, 2022

View reviewed changes

src/datasets/iterable_dataset.py Outdated Show resolved Hide resolved

src/datasets/utils/patching.py Outdated Show resolved Hide resolved

lhoestq mentioned this pull request Jun 9, 2022

Extend support for streaming datasets that use xml.dom.minidom.parse #4464

Merged

lhoestq added 2 commits June 9, 2022 11:10

comments + raise error if sharding is ambiguous

8f5579e

Merge branch 'master' into parallel-torch-iterable-dataset

ab91dbd

minor

1b87fb3

Merge branch 'master' into parallel-torch-iterable-dataset

b675a69

mariosasko approved these changes Jun 10, 2022

View reviewed changes

src/datasets/iterable_dataset.py Outdated Show resolved Hide resolved

lhoestq and others added 2 commits June 10, 2022 18:37

Update src/datasets/iterable_dataset.py

816d591

Co-authored-by: Mario Šaško <mariosasko777@gmail.com>

Merge branch 'master' into parallel-torch-iterable-dataset

ff586c4

lhoestq merged commit ab7d304 into master Jun 10, 2022

lhoestq deleted the parallel-torch-iterable-dataset branch June 10, 2022 20:47

albertvillanova mentioned this pull request Jun 15, 2022

Fix patching module that doesn't exist #4495

Merged

abhi-mosaic mentioned this pull request Jul 1, 2022

BERT and C4 updates mosaicml/composer#1252

Merged

lhoestq mentioned this pull request Jul 4, 2022

Streaming with num_workers != 0 #3148

Closed

lhoestq mentioned this pull request Jul 6, 2022

data duplicate when setting num_works > 1 with streaming data #3423

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Support DataLoader with num_workers > 0 in streaming mode #4375

Support DataLoader with num_workers > 0 in streaming mode #4375

lhoestq commented May 19, 2022 •

edited

Loading

HuggingFaceDocBuilderDev commented May 19, 2022 •

edited

Loading

lhoestq commented Jun 7, 2022

lhoestq Jun 7, 2022

lhoestq Jun 7, 2022

lhoestq Jun 7, 2022

lhoestq Jun 7, 2022 •

edited

Loading

mariosasko left a comment

lhoestq commented Jun 9, 2022

mariosasko commented Jun 9, 2022

lhoestq commented Jun 10, 2022 •

edited

Loading

mariosasko left a comment

lhoestq commented Jun 10, 2022

justheuristic commented Jul 4, 2022

Support DataLoader with num_workers > 0 in streaming mode #4375

Support DataLoader with num_workers > 0 in streaming mode #4375

Conversation

lhoestq commented May 19, 2022 • edited Loading

Issue

Solution in this PR

A few details regarding fsspec in multiprocessing

HuggingFaceDocBuilderDev commented May 19, 2022 • edited Loading

lhoestq commented Jun 7, 2022

lhoestq Jun 7, 2022

Choose a reason for hiding this comment

lhoestq Jun 7, 2022

Choose a reason for hiding this comment

lhoestq Jun 7, 2022

Choose a reason for hiding this comment

lhoestq Jun 7, 2022 • edited Loading

Choose a reason for hiding this comment

mariosasko left a comment

Choose a reason for hiding this comment

lhoestq commented Jun 9, 2022

mariosasko commented Jun 9, 2022

lhoestq commented Jun 10, 2022 • edited Loading

mariosasko left a comment

Choose a reason for hiding this comment

lhoestq commented Jun 10, 2022

justheuristic commented Jul 4, 2022

lhoestq commented May 19, 2022 •

edited

Loading

A few details regarding `fsspec` in multiprocessing

HuggingFaceDocBuilderDev commented May 19, 2022 •

edited

Loading

lhoestq Jun 7, 2022 •

edited

Loading

lhoestq commented Jun 10, 2022 •

edited

Loading