[DataLoader] Close open in DataPipe streams on best effort basis #78952

VitalyFedyunin · 2022-06-06T18:28:20Z

Adding ability to:

Track open StreamWrappers with StreamWrapper.session_streams
Automatically close parent StreamWrapper (ex. torchdata tar is the parent and extracted file streams are children)
Close streams in discarded by filtering structures

Stack from ghstack (oldest at bottom):

-> [DataLoader] Close open in DataPipe streams on best effort basis #78952

Differential Revision: D37489935

[ghstack-poisoned]

facebook-github-bot · 2022-06-06T18:28:25Z

🔗 Helpful links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/78952
📄 Preview Python docs built from this PR
📄 Preview C++ docs built from this PR
❓Need help or want to give feedback on the CI? Visit our office hours

✅ No Failures (0 Pending)

As of commit c9c6bf8 (more details on the Dr. CI page):

Expand to see more

💚 💚 Looks good so far! There are no failures yet. 💚 💚

This comment was automatically generated by Dr. CI (expand for details).

Please report bugs/suggestions to the (internal) Dr. CI Users group.

Click here to manually regenerate this comment.

[ghstack-poisoned]

… basis" Adding ability to: - Track open StreamWrappers with `StreamWrapper.session_streams` - Automatically close parent StreamWrapper (ex. torchdata tar is the parent and extracted file streams are children) - Close streams in discarded by filtering structures [ghstack-poisoned]

NivekT

What happens if some does fork then filter (or zip)? Would this close the stream held by the other ChildDataPipe?

stream_dp = ...
cdp1, cdp2 = stream_dp.fork(2)
cdp2 = cdp2.filter(...)
list(cdp1.readlines())

I think we have to increment the counter when fork is used, and only close a stream when there is no other reference to it.

VitalyFedyunin · 2022-06-06T20:27:46Z

What happens if some does fork then filter (or zip)? Would this close the stream held by the other ChildDataPipe?
stream_dp = ...
cdp1, cdp2 = stream_dp.fork(2)
cdp2 = cdp2.filter(...)
list(cdp1.readlines())
I think we have to increment the counter when fork is used, and only close a stream when there is no other reference to it.

Ideally, we should warn and RuntimeError if stream is forked. IMO there is no good way to allow reading abstract buffer twice.

… basis" Adding ability to: - Track open StreamWrappers with `StreamWrapper.session_streams` - Automatically close parent StreamWrapper (ex. torchdata tar is the parent and extracted file streams are children) - Close streams in discarded by filtering structures [ghstack-poisoned]

ghstack-source-id: 3a5317f Pull Request resolved: #78952

ejguan · 2022-06-06T20:35:57Z

torch/utils/data/datapipes/iter/combining.py

+        iterators = [iter(datapipe) for datapipe in self.datapipes]
+        for data in zip(*iterators):
            yield data

+        unused = []
+        for iterator in iterators:
+            unused += list(iterator)
+
+        # TODO(VitalyFedyunin): This should be Exception or warning when torchdata.debug is enabled
+        for item in unused:
+            StreamWrapper.close_streams(item)


I am not sure about the behavior of generator functions. When a not fully consumed iterator is deconstructed, I think the iterator would be destroyed and these StreamWrappers are never created.

Do we need to handle this case?

I think there are two options here:

Do try...finally within __iter__, where the finally clause does the clean up.

This should work regardless how/why the deconstruction of the iterator happens.
I just tried this with IterableWrapper.

def __iter__(self): try: source_data = self.iterable if self.deepcopy: try: source_data = copy.deepcopy(self.iterable) except TypeError: warnings.warn("....") for data in source_data: yield data finally: print("FINALLLY")

dp = IterableWrapper(range(10)) it = iter(dp) print(next(it)) print("About to create a new iterator") it = iter(dp) print("Created a new iterator")

0 About to create a new iterator FINALLLY Created a new iterator

We can put the clean up logic with hook_iterator or reset method of IterDataPipe whenever that is triggered.

I like the idea of finally: it will produce readable code and can be placed as a development pattern

ejguan · 2022-06-06T20:39:15Z

torch/utils/data/datapipes/utils/common.py

@@ -157,13 +157,59 @@ class StreamWrapper:
    DataPipe operation like `FileOpener`. StreamWrapper would guarantee
    the wrapped file handler is closed when it's out of scope.
    '''
-    def __init__(self, file_obj):
+    session_streams: Dict[Any, int] = {}


We can use a set in this case.
Do you think we might also register a function to atexit to make sure all session_streams are closed?

Oddly it doesn't allow to put self into Set. Switching back to Dict.

…d streams" Blocked by pytorch/pytorch#78952 Automatically cleans StreamsWrappers from various buffers Automatically closing parent unarchive streams [ghstack-poisoned]

… basis" Adding ability to: - Track open StreamWrappers with `StreamWrapper.session_streams` - Automatically close parent StreamWrapper (ex. torchdata tar is the parent and extracted file streams are children) - Close streams in discarded by filtering structures [ghstack-poisoned]

ghstack-source-id: 251e5cf Pull Request resolved: #78952

… basis" Adding ability to: - Track open StreamWrappers with `StreamWrapper.session_streams` - Automatically close parent StreamWrapper (ex. torchdata tar is the parent and extracted file streams are children) - Close streams in discarded by filtering structures [ghstack-poisoned]

ghstack-source-id: df41ef1 Pull Request resolved: #78952

VitalyFedyunin · 2022-06-28T15:56:43Z

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

NivekT

I think we need some tests to reason through the common and complicated cases.
Question: when a buffer is cleared (say in Shuffler), do we need to attempt to call close_streams?

NivekT · 2022-06-28T18:51:11Z

torch/utils/data/datapipes/utils/common.py

+        self.name = name
+        if parent_stream is not None:
+            if not isinstance(parent_stream, StreamWrapper):
+                raise RuntimeError('Parent steam should be StreamWrapper, {} was given'.format(type(parent_stream)))


Suggested change

raise RuntimeError('Parent steam should be StreamWrapper, {} was given'.format(type(parent_stream)))

raise RuntimeError('Parent stream should be StreamWrapper, {} was given'.format(type(parent_stream)))

NivekT · 2022-06-28T18:57:11Z

torch/utils/data/datapipes/utils/common.py

+    def autoclose(self):
+        '''
+        Marks Steam to close automatically as soon as all child streams are closed.
+        '''
+        if self.child_counter == 0:
+            self.close()
+        self.close_on_last_child = True


What will be the usage of this? Provide an option for mark streams as autoclose?
If that is the case, maybe that can be an optional argument to __init__?

It will close stream if there is no active children or mark it to be closed automatically if there is at least one.

NivekT · 2022-06-28T19:05:46Z

torch/utils/data/datapipes/utils/common.py

+        if self.child_counter == 0:
+            self.close()
+        self.close_on_last_child = True
+


Do we need to change the implementation of __del__ here to remove self from StreamWrapper.session_streams?

I agree on this comment

ejguan · 2022-06-28T19:31:20Z

torch/utils/data/datapipes/utils/common.py

+            if isinstance(v, dict):
+                for kk, vv in v.items():
+                    cls.close_streams(vv, depth=depth + 1)
+            elif isinstance(v, list) or isinstance(v, tuple):


nit: isinstance(v, (list, tuple))

ejguan · 2022-06-28T19:36:56Z

torch/utils/data/datapipes/utils/common.py

+            self.parent_stream.child_counter -= 1
+            if not self.parent_stream.child_counter and self.parent_stream.close_on_last_child:
+                self.parent_stream.close()
+        self.file_obj.close(*args, **kwargs)


Suggested change

self.file_obj.close(*args, **kwargs)

try:

self.file_obj.close(*args, **kwargs)

except AttributeError:

pass

… basis" Adding ability to: - Track open StreamWrappers with `StreamWrapper.session_streams` - Automatically close parent StreamWrapper (ex. torchdata tar is the parent and extracted file streams are children) - Close streams in discarded by filtering structures Differential Revision: [D37489935](https://our.internmc.facebook.com/intern/diff/D37489935) [ghstack-poisoned]

ghstack-source-id: d77d7c0 Pull Request resolved: #78952

VitalyFedyunin · 2022-06-29T17:05:47Z

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

VitalyFedyunin · 2022-06-29T17:54:14Z

@pytorchbot merge -g

pytorchmergebot · 2022-06-29T17:55:45Z

@pytorchbot successfully started a merge job. Check the current status here

pytorchmergebot · 2022-06-29T20:11:25Z

@VitalyFedyunin your PR has been successfully merged.

github-actions · 2022-06-29T20:12:05Z

Hey @VitalyFedyunin.
You've committed this PR, but it does not have both a 'release notes: ...' and 'topics: ...' label. Please add one of each to the PR. The 'release notes: ...' label should represent the part of PyTorch that this PR changes (fx, autograd, distributed, etc) and the 'topics: ...' label should represent the kind of PR it is (not user facing, new feature, bug fix, perf improvement, etc). The list of valid labels can be found here for the 'release notes: ...' and here for the 'topics: ...'.
For changes that are 'topic: not user facing' there is no need for a release notes label.

) (#78952) Summary: Adding ability to: - Track open StreamWrappers with `StreamWrapper.session_streams` - Automatically close parent StreamWrapper (ex. torchdata tar is the parent and extracted file streams are children) - Close streams in discarded by filtering structures Pull Request resolved: #78952 Approved by: https://github.com/ejguan Test Plan: contbuild & OSS CI, see https://hud.pytorch.org/commit/pytorch/pytorch/331c0c18033bf4138c9ea468a8c759865dd8ffff Reviewed By: ejguan Differential Revision: D37489935 fbshipit-source-id: 6c0fa4d03b1d957cae9ab6062a00b27b120d68e6

Debug pipes overflow

e7026bd

[ghstack-poisoned]

facebook-github-bot added the cla signed label Jun 6, 2022

Update on "Debug pipes overflow"

cff8029

[ghstack-poisoned]

Update on "Debug pipes overflow"

6ff14b8

[ghstack-poisoned]

VitalyFedyunin changed the title ~~Debug pipes overflow~~ [DataLoader] Close open in DataPipe streams on best effort basis Jun 6, 2022

VitalyFedyunin requested review from NivekT and ejguan June 6, 2022 18:42

NivekT reviewed Jun 6, 2022

View reviewed changes

VitalyFedyunin added a commit that referenced this pull request Jun 6, 2022

Debug pipes overflow

3ede7dc

ghstack-source-id: 3a5317f Pull Request resolved: #78952

ejguan reviewed Jun 6, 2022

View reviewed changes

VitalyFedyunin mentioned this pull request Jun 6, 2022

[DataPipe] Automatically close parent streams and discarded streams pytorch/data#498

Closed

VitalyFedyunin added a commit that referenced this pull request Jun 27, 2022

Debug pipes overflow

f8f2e3f

ghstack-source-id: 251e5cf Pull Request resolved: #78952

VitalyFedyunin requested review from ejguan and NivekT June 27, 2022 21:57

VitalyFedyunin added a commit that referenced this pull request Jun 28, 2022

Debug pipes overflow

f6be168

ghstack-source-id: df41ef1 Pull Request resolved: #78952

NivekT reviewed Jun 28, 2022

View reviewed changes

ejguan approved these changes Jun 28, 2022

View reviewed changes

VitalyFedyunin added a commit that referenced this pull request Jun 29, 2022

Debug pipes overflow

5006155

ghstack-source-id: d77d7c0 Pull Request resolved: #78952

pytorchmergebot added the Merged label Jun 29, 2022

pytorchmergebot closed this in 331c0c1 Jun 29, 2022

NivekT added module: data torch.utils.data release notes: dataloader release notes category topic: improvements topic category labels Jun 29, 2022

facebook-github-bot deleted the gh/VitalyFedyunin/179/head branch July 3, 2022 14:17

VitalyFedyunin mentioned this pull request Jul 6, 2022

Closing streams to avoid testing issues pytorch/vision#6128

Closed

	raise RuntimeError('Parent steam should be StreamWrapper, {} was given'.format(type(parent_stream)))
	raise RuntimeError('Parent stream should be StreamWrapper, {} was given'.format(type(parent_stream)))

[DataLoader] Close open in DataPipe streams on best effort basis #78952

[DataLoader] Close open in DataPipe streams on best effort basis #78952

Uh oh!

Conversation

VitalyFedyunin commented Jun 6, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

facebook-github-bot commented Jun 6, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful links

✅ No Failures (0 Pending)

Uh oh!

NivekT left a comment

Choose a reason for hiding this comment

Uh oh!

VitalyFedyunin commented Jun 6, 2022

Uh oh!

ejguan Jun 6, 2022 • 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!

NivekT Jun 6, 2022 • 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!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

VitalyFedyunin commented Jun 28, 2022

Uh oh!

NivekT left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

VitalyFedyunin commented Jun 29, 2022

Uh oh!

VitalyFedyunin commented Jun 29, 2022

Uh oh!

pytorchmergebot commented Jun 29, 2022

Uh oh!

pytorchmergebot commented Jun 29, 2022

Uh oh!

github-actions bot commented Jun 29, 2022

Uh oh!

Uh oh!

VitalyFedyunin commented Jun 6, 2022 •

edited

Loading

facebook-github-bot commented Jun 6, 2022 •

edited

Loading

ejguan Jun 6, 2022 •

edited

Loading

NivekT Jun 6, 2022 •

edited

Loading