Debug pipes overflow

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

Author: VitalyFedyunin

torch/utils/data/datapipes/datapipe.pyi.in

@@ -4,7 +4,7 @@
 # classes/objects here, even though we are not injecting extra code into them at the moment.
 
 from torch.utils.data.datapipes._typing import _DataPipeMeta, _IterDataPipeMeta
-from typing import Any, Callable, Dict, Generic, Iterator, List, Optional, TypeVar
+from typing import Any, Callable, Dict, Generic, Iterator, List, Optional, TypeVar, Union
 from torch.utils.data import Dataset, IterableDataset
 
 T_co = TypeVar('T_co', covariant=True)

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

@@ -5,7 +5,7 @@
 
 from torch.utils.data.datapipes._decorator import functional_datapipe
 from torch.utils.data.datapipes.datapipe import IterDataPipe
-from torch.utils.data.datapipes.utils.common import _check_lambda_fn
+from torch.utils.data.datapipes.utils.common import StreamWrapper, _check_lambda_fn
 
 __all__ = [
     "ConcaterIterDataPipe",
@@ -345,6 +345,7 @@ def _find_next(self, instance_id: int) -> T_co:
             value = next(self._datapipe_iterator)
             classification = self.classifier_fn(value)
             if classification is None and self.drop_none:
+                StreamWrapper.close_streams(value)
                 continue
             if classification is None or classification >= self.num_instances or classification < 0:
                 raise ValueError(f"Output of the classification fn should be between 0 and {self.num_instances - 1}. " +
@@ -510,8 +511,18 @@ def __init__(self, *datapipes: IterDataPipe):
         self.length = None
 
     def __iter__(self) -> Iterator[Tuple[T_co]]:
-        for data in zip(*self.datapipes):
-            yield data
+        iterators = [iter(datapipe) for datapipe in self.datapipes]
+        try:
+            for data in zip(*iterators):
+                yield data
+        finally:
+            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)
 
     def __len__(self) -> int:
         if self.length is not None:

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

@@ -4,6 +4,8 @@
 from torch.utils.data.datapipes.datapipe import IterDataPipe
 from torch.utils.data.datapipes.dataframe import dataframe_wrapper as df_wrapper
 from torch.utils.data.datapipes.utils.common import _check_lambda_fn, _deprecation_warning
+from torch.utils.data.datapipes.utils.common import StreamWrapper
+
 
 __all__ = ["FilterIterDataPipe", ]
 
@@ -78,6 +80,8 @@ def __iter__(self) -> Iterator[T_co]:
             filtered = self._returnIfTrue(data)
             if self._isNonEmpty(filtered):
                 yield filtered
+            else:
+                StreamWrapper.close_streams(data)
 
     def _returnIfTrue(self, data):
         condition = self._apply_filter_fn(data)

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

@@ -32,5 +32,6 @@ def __iter__(self):
             while True:
                 d = stream.read(self.chunk)
                 if not d:
+                    stream.close()
                     break
                 yield (furl, d)

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

@@ -3,7 +3,8 @@
 import warnings
 
 from io import IOBase
-from typing import Dict, Iterable, List, Tuple, Union, Optional
+from typing import Any, Dict, Iterable, List, Set, Tuple, Union, Optional
+
 
 from torch.utils.data._utils.serialization import DILL_AVAILABLE
 
@@ -179,13 +180,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: Set[Any] = set()
+
+    def __init__(self, file_obj, parent_stream=None, name=None):
         self.file_obj = file_obj
+        self.child_counter = 0
+        self.parent_stream = parent_stream
+        self.close_on_last_child = False
+        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)))
+            parent_stream.child_counter += 1
+            self.parent_stream = parent_stream
+        StreamWrapper.session_streams.update(self)
+
+    @classmethod
+    def close_streams(cls, v, depth=0):
+        '''
+        Traverse structure and attempts to close all found StreamWrappers on best effort basis.
+        '''
+        if depth > 10:
+            return
+        if isinstance(v, StreamWrapper):
+            v.close()
+        else:
+            # Traverve only simple structures
+            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):
+                for vv in v:
+                    cls.close_streams(vv, depth=depth + 1)
 
     def __getattr__(self, name):
         file_obj = self.__dict__['file_obj']
         return getattr(file_obj, name)
 
+    def close(self, *args, **kwargs):
+        StreamWrapper.session_streams.remove(self)
+        if self.parent_stream is not None:
+            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)
+
+    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
+
     def __dir__(self):
         attrs = list(self.__dict__.keys()) + list(StreamWrapper.__dict__.keys())
         attrs += dir(self.file_obj)
@@ -205,7 +252,10 @@ def __next__(self):
         return next(self.file_obj)
 
     def __repr__(self):
-        return f"StreamWrapper<{self.file_obj!r}>"
+        if self.name is None:
+            return f"StreamWrapper<{self.file_obj!r}>"
+        else:
+            return f"StreamWrapper<{self.name},{self.file_obj!r}>"
 
     def __getstate__(self):
         return self.file_obj