[PrototypeRS] Adding 'pause' and 'resume' operations to halt DataPipes by NivekT · Pull Request #879 · meta-pytorch/data

NivekT · 2022-11-03T22:44:44Z

Stack from ghstack:

The goal of this PR is fully stop DataPipe activities in preparation of snapshotting (which requires a halted state), so buffers will not be changing while we record the snapshot.

This PR doesn't handle pausing or resuming with sharding_round_robin_dispatch. This will be fixed in an upcoming PR and is tracked by #991.

Implementation note (feel free to discuss):
_pause exists for both DataLoader2Iterator and DataLoader2 as private methods.
resume exists for both DataLoader2Iterator as public but it is private for DataLoader2.

limit and clear_limit exist only for DataLoader2Iterator as a public methods.

Limit persists through iterations (with the iterator) until it is manually cleared via clear_limit.
Creating a new iterator (iter(dl2)) will not retain the previous limit, because the limit is attached to the iterator.

Differential Revision: D41744759

[ghstack-poisoned]

ghstack-source-id: 75446ff Pull Request resolved: #879

NivekT · 2022-11-03T22:46:23Z

torchdata/dataloader2/communication/iter.py

+            graph = traverse_dps(source_datapipe)
+            for dp, _ in graph.values():
+                if hasattr(dp, "full_stop") and callable(dp.full_stop):
+                    dp.full_stop()
+            protocol.response_full_stop()


Is there any potential issue with traversing through the graph and calling .full_stop() on all of them?

torchdata/dataloader2/reading_service.py

torchdata/datapipes/iter/util/prefetcher.py

NivekT · 2022-11-04T18:45:41Z

Another thing to discuss: should we do double underscore or something to protect __reset`` and __full_stop` from users?

The goal of this PR is fully stop DataPipe activities in preparation of snapshotting (which requires a halted state), so buffers will not be changing while we record the snapshot. [ghstack-poisoned]

ghstack-source-id: f707439 Pull Request resolved: #879

NivekT

Updated the implementation but I still need some test cases. Suggestions are welcomed.

torchdata/datapipes/iter/util/prefetcher.py

ejguan

Can we change the name to pause? full stop seems like an operation ran at the end epoch

torchdata/datapipes/iter/util/prefetcher.py

The goal of this PR is fully stop DataPipe activities in preparation of snapshotting (which requires a halted state), so buffers will not be changing while we record the snapshot. [ghstack-poisoned]

ghstack-source-id: 15842d2 Pull Request resolved: #879

The goal of this PR is fully stop DataPipe activities in preparation of snapshotting (which requires a halted state), so buffers will not be changing while we record the snapshot. [ghstack-poisoned]

ghstack-source-id: 9184795 Pull Request resolved: #879

…lt DataPipes" The goal of this PR is fully stop DataPipe activities in preparation of snapshotting (which requires a halted state), so buffers will not be changing while we record the snapshot. [ghstack-poisoned]

ghstack-source-id: 9164020 Pull Request resolved: #879

…lt DataPipes" The goal of this PR is fully stop DataPipe activities in preparation of snapshotting (which requires a halted state), so buffers will not be changing while we record the snapshot. [ghstack-poisoned]

ghstack-source-id: a99c23a Pull Request resolved: #879

…lt DataPipes" The goal of this PR is fully stop DataPipe activities in preparation of snapshotting (which requires a halted state), so buffers will not be changing while we record the snapshot. [ghstack-poisoned]

ghstack-source-id: 023aa69 Pull Request resolved: #879

…lt DataPipes" The goal of this PR is fully stop DataPipe activities in preparation of snapshotting (which requires a halted state), so buffers will not be changing while we record the snapshot. [ghstack-poisoned]

ghstack-source-id: dfb19de Pull Request resolved: #879

ejguan · 2022-11-17T14:11:38Z

I forget to mention about one scenario. If I iterate datapipe after pause but before resume, I expect a StopIteration is raised. Am I right?

>>> dp.pause()
>>> for d in dp:
...     print(d)  # Nothing should be printed
>>> dp.resume()
>>> for d in dp:
...     print(d)
# Values

NivekT · 2022-11-17T15:30:44Z

I forget to mention about one scenario. If I iterate datapipe after pause but before resume, I expect a StopIteration is raised. Am I right?
>>> dp.pause()
>>> for d in dp:
...     print(d)  # Nothing should be printed
>>> dp.resume()
>>> for d in dp:
...     print(d)
# Values

If you try to call next after pause, this implementation raises RuntimeError.

ejguan · 2022-11-17T15:36:36Z

If you try to call next after pause, this implementation raises RuntimeError.

If we want to achieve mini-epoch for OSS as well, we might need it to be StopIteration.

NivekT · 2022-11-17T15:48:53Z

I see. I think what we can do is:

Inside the internal of PrototypeRS, if it cannot submit a GetNextRequest after pause without calling resume first. This is implemented in this PR.

Publicly, DataLoader2 can behave different (maybe like you said return StopIteration), but it needs to do so without triggering reset when a new iterator is made? That might get a bit tricky and I need to think more about it. This part is not covered by this PR (since I'm not exposing public API of pause or resume here).

torchdata/dataloader2/communication/iter.py

…lt DataPipes" The goal of this PR is fully stop DataPipe activities in preparation of snapshotting (which requires a halted state), so buffers will not be changing while we record the snapshot. [ghstack-poisoned]

test/dataloader2/test_dataloader2.py

…lt DataPipes" The goal of this PR is fully stop DataPipe activities in preparation of snapshotting (which requires a halted state), so buffers will not be changing while we record the snapshot. This PR doesn't handle pausing or resuming with `sharding_round_robin_dispatch`. This will be fixed in an upcoming PR and is tracked by #991. --- Implementation note (feel free to discuss): `_pause` exists for both `DataLoader2Iterator` and `DataLoader2` as private methods. `resume` exists for both `DataLoader2Iterator` as public but it is private for `DataLoader2`. `limit` and `clear_limit` exist only for `DataLoader2Iterator` as a public methods. - Limit persists through iterations (with the iterator) until it is manually cleared via `clear_limit`. - Creating a new iterator (`iter(dl2)`) will not retain the previous limit, because the limit is attached to the iterator. Differential Revision: [D41744759](https://our.internmc.facebook.com/intern/diff/D41744759) [ghstack-poisoned]

torchdata/dataloader2/reading_service.py

ejguan

The PR looks good in terms of test cases. I left a few comments below

torchdata/dataloader2/reading_service.py

ejguan · 2023-02-07T20:56:36Z

torchdata/datapipes/iter/util/distributed.py

+            self._executor = None
+
+    def resume(self):
+        self._executor = _PrefetchExecutor(iter(self.datapipe), 1, self._callback_fn, self.timeout)


It seems that the resume for FullSync is going to re-read from the beginning. Is this expected? Or, we can raise Exception for fullsync saying it doesn't support checkpointing yet.

Is there a way for it not to read from the beginning again? Halting the executor by changing __iter__?

For now, I can raise an Exception in both pause and resume.

torchdata/dataloader2/communication/protocol.py

torchdata/dataloader2/communication/iter.py

test/dataloader2/test_proto_multi_rs.py

NivekT · 2023-02-07T21:30:08Z

CI is expected to fail until PyTorch core nightly updates

…lt DataPipes" The goal of this PR is fully stop DataPipe activities in preparation of snapshotting (which requires a halted state), so buffers will not be changing while we record the snapshot. This PR doesn't handle pausing or resuming with `sharding_round_robin_dispatch`. This will be fixed in an upcoming PR and is tracked by #991. --- Implementation note (feel free to discuss): `_pause` exists for both `DataLoader2Iterator` and `DataLoader2` as private methods. `resume` exists for both `DataLoader2Iterator` as public but it is private for `DataLoader2`. `limit` and `clear_limit` exist only for `DataLoader2Iterator` as a public methods. - Limit persists through iterations (with the iterator) until it is manually cleared via `clear_limit`. - Creating a new iterator (`iter(dl2)`) will not retain the previous limit, because the limit is attached to the iterator. Differential Revision: [D41744759](https://our.internmc.facebook.com/intern/diff/D41744759) [ghstack-poisoned]

NivekT · 2023-02-08T01:05:39Z

@NivekT has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

…lt DataPipes" The goal of this PR is fully stop DataPipe activities in preparation of snapshotting (which requires a halted state), so buffers will not be changing while we record the snapshot. This PR doesn't handle pausing or resuming with `sharding_round_robin_dispatch`. This will be fixed in an upcoming PR and is tracked by #991. --- Implementation note (feel free to discuss): `_pause` exists for both `DataLoader2Iterator` and `DataLoader2` as private methods. `resume` exists for both `DataLoader2Iterator` as public but it is private for `DataLoader2`. `limit` and `clear_limit` exist only for `DataLoader2Iterator` as a public methods. - Limit persists through iterations (with the iterator) until it is manually cleared via `clear_limit`. - Creating a new iterator (`iter(dl2)`) will not retain the previous limit, because the limit is attached to the iterator. Differential Revision: [D41744759](https://our.internmc.facebook.com/intern/diff/D41744759) [ghstack-poisoned]

NivekT · 2023-02-08T17:58:52Z

Addressed and resolved most of the comments but 2. CI is being re-ran after nightly updated.

ejguan

Overall LGTM. Thank you. I left a few comments to help simplify the test and a comment for internal users

test/dataloader2/test_proto_multi_rs.py

torchdata/dataloader2/dataloader2.py

…lt DataPipes" The goal of this PR is fully stop DataPipe activities in preparation of snapshotting (which requires a halted state), so buffers will not be changing while we record the snapshot. This PR doesn't handle pausing or resuming with `sharding_round_robin_dispatch`. This will be fixed in an upcoming PR and is tracked by #991. --- Implementation note (feel free to discuss): `_pause` exists for both `DataLoader2Iterator` and `DataLoader2` as private methods. `resume` exists for both `DataLoader2Iterator` as public but it is private for `DataLoader2`. `limit` and `clear_limit` exist only for `DataLoader2Iterator` as a public methods. - Limit persists through iterations (with the iterator) until it is manually cleared via `clear_limit`. - Creating a new iterator (`iter(dl2)`) will not retain the previous limit, because the limit is attached to the iterator. Differential Revision: [D41744759](https://our.internmc.facebook.com/intern/diff/D41744759) [ghstack-poisoned]

NivekT · 2023-02-08T19:53:37Z

@NivekT has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

…lt DataPipes" The goal of this PR is fully stop DataPipe activities in preparation of snapshotting (which requires a halted state), so buffers will not be changing while we record the snapshot. This PR doesn't handle pausing or resuming with `sharding_round_robin_dispatch`. This will be fixed in an upcoming PR and is tracked by #991. --- Implementation note (feel free to discuss): `_pause` exists for both `DataLoader2Iterator` and `DataLoader2` as private methods. `resume` exists for both `DataLoader2Iterator` as public but it is private for `DataLoader2`. `limit` and `clear_limit` exist only for `DataLoader2Iterator` as a public methods. - Limit persists through iterations (with the iterator) until it is manually cleared via `clear_limit`. - Creating a new iterator (`iter(dl2)`) will not retain the previous limit, because the limit is attached to the iterator. Differential Revision: [D41744759](https://our.internmc.facebook.com/intern/diff/D41744759) [ghstack-poisoned]

NivekT · 2023-02-09T23:02:17Z

@NivekT has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

…lt DataPipes" The goal of this PR is fully stop DataPipe activities in preparation of snapshotting (which requires a halted state), so buffers will not be changing while we record the snapshot. This PR doesn't handle pausing or resuming with `sharding_round_robin_dispatch`. This will be fixed in an upcoming PR and is tracked by #991. --- Implementation note (feel free to discuss): `_pause` exists for both `DataLoader2Iterator` and `DataLoader2` as private methods. `resume` exists for both `DataLoader2Iterator` as public but it is private for `DataLoader2`. `limit` and `clear_limit` exist only for `DataLoader2Iterator` as a public methods. - Limit persists through iterations (with the iterator) until it is manually cleared via `clear_limit`. - Creating a new iterator (`iter(dl2)`) will not retain the previous limit, because the limit is attached to the iterator. Differential Revision: [D41744759](https://our.internmc.facebook.com/intern/diff/D41744759) [ghstack-poisoned]

NivekT · 2023-02-10T21:58:25Z

@NivekT has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

NivekT · 2023-02-13T21:48:15Z

@NivekT has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

facebook-github-bot · 2023-02-14T04:19:06Z

@NivekT merged this pull request in cfcdf6e.

ejguan · 2023-02-14T23:11:31Z

torchdata/datapipes/iter/util/prefetcher.py

+            self.thread = thread
+            self.thread.start()
+
+            while prefetch_data.run_prefetcher:


One more thing: Should this line be while not prefetch_data.stop_iteration:? The current implementation means that iterator exits earlier than it's supposed to be.

It's a little bit fuzzy to me if we are able to call next over the iterator of this DataPipe after pause.

I guess my main question is if pause of DataPipe is exposed to users.

I think you are right. Though from a glance changing that while condition will also require changing the exception conditions.

Particularly, it should become except StopIteration: prefetch_data.run_prefetcher = False and the same for except communication.iter.InvalidStateResetRequired.

Overall, it should be:
prefetch_data.run_prefetcher - indicates whether or not we should be prefetching, it should become False when the source runs out of elements, invalid state, termination, and maybe other exception raised.
prefetch_data.stop_iteration - this should. indicate when we should stop yielding elements, it should mostly be the same as above except when the source is out of elements, it should keep yielding until the buffer is empty.

There maybe other corner cases I have not stated here. We can test it out and see.

ejguan · 2023-02-14T23:19:14Z

torchdata/datapipes/iter/util/prefetcher.py

+                    except StopIteration:
+                        prefetch_data.stop_iteration = True
+                    except communication.iter.InvalidStateResetRequired:
+                        prefetch_data.stop_iteration = True


Related to my last comment. stop_iteration should dominate run_prefetcher. When stop_iterator is met, there is no reason to process with prefetching, right?

Whenever stop_iteration turns to True, run_prefetcher should be False

What you stated should be True, but it requires changing other parts as well.

[PrototypeRS] Adding 'full stop' to halt DataPipes

d22790b

[ghstack-poisoned]

NivekT mentioned this pull request Nov 3, 2022

[ReadingService] Small clean up and comments for PrototypeRS #878

Closed

NivekT added a commit that referenced this pull request Nov 3, 2022

[PrototypeRS] Adding 'full stop' to halt DataPipes

2c3f260

ghstack-source-id: 75446ff Pull Request resolved: #879

facebook-github-bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Nov 3, 2022

NivekT marked this pull request as draft November 3, 2022 22:45

NivekT commented Nov 4, 2022

View reviewed changes

Update on "[PrototypeRS] Adding 'full stop' to halt DataPipes"

9efb6b3

The goal of this PR is fully stop DataPipe activities in preparation of snapshotting (which requires a halted state), so buffers will not be changing while we record the snapshot. [ghstack-poisoned]

NivekT added a commit that referenced this pull request Nov 7, 2022

[PrototypeRS] Adding 'full stop' to halt DataPipes

06c80b6

ghstack-source-id: f707439 Pull Request resolved: #879

NivekT commented Nov 7, 2022

View reviewed changes

torchdata/datapipes/iter/util/prefetcher.py Outdated Show resolved Hide resolved

ejguan reviewed Nov 8, 2022

View reviewed changes

torchdata/datapipes/iter/util/prefetcher.py Outdated Show resolved Hide resolved

torchdata/datapipes/iter/util/prefetcher.py Outdated Show resolved Hide resolved

Update on "[PrototypeRS] Adding 'full stop' to halt DataPipes"

50bea6c

The goal of this PR is fully stop DataPipe activities in preparation of snapshotting (which requires a halted state), so buffers will not be changing while we record the snapshot. [ghstack-poisoned]

NivekT added a commit that referenced this pull request Nov 10, 2022

[PrototypeRS] Adding 'pause' and 'resume' operations to halt DataPipes

f80b57c

ghstack-source-id: 15842d2 Pull Request resolved: #879

Update on "[PrototypeRS] Adding 'full stop' to halt DataPipes"

a853244

The goal of this PR is fully stop DataPipe activities in preparation of snapshotting (which requires a halted state), so buffers will not be changing while we record the snapshot. [ghstack-poisoned]

NivekT added a commit that referenced this pull request Nov 11, 2022

[PrototypeRS] Adding 'pause' and 'resume' operations to halt DataPipes

31baf76

ghstack-source-id: 9184795 Pull Request resolved: #879

NivekT changed the title ~~[PrototypeRS] Adding 'full stop' to halt DataPipes~~ [PrototypeRS] Adding 'pause' and 'resume' operations to halt DataPipes Nov 11, 2022

NivekT added a commit that referenced this pull request Nov 14, 2022

[PrototypeRS] Adding 'pause' and 'resume' operations to halt DataPipes

41a9215

ghstack-source-id: 9164020 Pull Request resolved: #879

NivekT added a commit that referenced this pull request Nov 16, 2022

[PrototypeRS] Adding 'pause' and 'resume' operations to halt DataPipes

2da9a50

ghstack-source-id: a99c23a Pull Request resolved: #879

NivekT added a commit that referenced this pull request Nov 16, 2022

[PrototypeRS] Adding 'pause' and 'resume' operations to halt DataPipes

78629be

ghstack-source-id: 023aa69 Pull Request resolved: #879

NivekT added a commit that referenced this pull request Nov 16, 2022

[PrototypeRS] Adding 'pause' and 'resume' operations to halt DataPipes

d87704f

ghstack-source-id: dfb19de Pull Request resolved: #879

ejguan reviewed Nov 17, 2022

View reviewed changes

torchdata/dataloader2/communication/iter.py Outdated Show resolved Hide resolved

NivekT commented Feb 6, 2023

View reviewed changes

test/dataloader2/test_dataloader2.py Show resolved Hide resolved

NivekT mentioned this pull request Feb 6, 2023

Enable checkpointing at the beginning of an epoch #992

Closed

ejguan reviewed Feb 7, 2023

View reviewed changes

torchdata/dataloader2/reading_service.py Outdated Show resolved Hide resolved

ejguan reviewed Feb 7, 2023

View reviewed changes

ejguan approved these changes Feb 8, 2023

View reviewed changes

NivekT mentioned this pull request Feb 8, 2023

[DataLoader2] Saving and restoring initial seed generator #998

Closed

NivekT added 2 commits February 9, 2023 17:47

NivekT mentioned this pull request Feb 13, 2023

[DataLoader2] Removing delegation for 'pause', 'limit', and 'resume' #1011

Closed

facebook-github-bot closed this in cfcdf6e Feb 14, 2023

facebook-github-bot added the Merged label Feb 14, 2023

ejguan reviewed Feb 14, 2023

View reviewed changes

NivekT mentioned this pull request Feb 15, 2023

Update Prefetcher and Implement PinMemory IterDataPipe #1014

Closed

facebook-github-bot deleted the gh/NivekT/100/head branch February 17, 2023 15:20

Conversation

NivekT commented Nov 3, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

NivekT Nov 3, 2022

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

NivekT commented Nov 4, 2022

Uh oh!

NivekT left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

ejguan left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

ejguan commented Nov 17, 2022

Uh oh!

NivekT commented Nov 17, 2022

Uh oh!

ejguan commented Nov 17, 2022

Uh oh!

NivekT commented Nov 17, 2022

Uh oh!

Uh oh!

Uh oh!

Uh oh!

ejguan left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

ejguan Feb 7, 2023

Choose a reason for hiding this comment

Uh oh!

NivekT Feb 7, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

NivekT commented Feb 7, 2023

Uh oh!

NivekT commented Feb 8, 2023

Uh oh!

NivekT commented Feb 8, 2023

Uh oh!

ejguan left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

NivekT commented Feb 8, 2023

Uh oh!

NivekT commented Feb 9, 2023

Uh oh!

NivekT commented Feb 10, 2023

Uh oh!

NivekT commented Feb 13, 2023

Uh oh!

facebook-github-bot commented Feb 14, 2023

Uh oh!

ejguan Feb 14, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ejguan Feb 14, 2023

Choose a reason for hiding this comment

NivekT commented Nov 3, 2022 •

edited

Loading

NivekT Feb 7, 2023 •

edited

Loading

ejguan Feb 14, 2023 •

edited

Loading

NivekT Feb 15, 2023 •

edited

Loading

ejguan Feb 14, 2023 •

edited

Loading