[DataPipe] Automatically close parent streams and discarded streams

ghstack-source-id: b0a21fa
Pull Request resolved: #560

Author: VitalyFedyunin

torchdata/__init__.py

@@ -8,9 +8,17 @@
 
 from . import datapipes
 
+janitor = datapipes.utils.janitor
+
 try:
     from .version import __version__  # noqa: F401
 except ImportError:
     pass
 
-__all__ = ["datapipes"]
+__all__ = [
+    "datapipes",
+    "janitor",
+]
+
+# Please keep this list sorted
+assert __all__ == sorted(__all__)

torchdata/datapipes/iter/util/bz2fileloader.py

@@ -54,10 +54,13 @@ def __iter__(self) -> Iterator[Tuple[str, BufferedIOBase]]:
             try:
                 extracted_fobj = bz2.open(data_stream, mode="rb")  # type: ignore[call-overload]
                 new_pathname = pathname.rstrip(".bz2")
-                yield new_pathname, StreamWrapper(extracted_fobj)  # type: ignore[misc]
+                yield new_pathname, StreamWrapper(extracted_fobj, data_stream, name=new_pathname)  # type: ignore[misc]
             except Exception as e:
                 warnings.warn(f"Unable to extract files from corrupted bzip2 stream {pathname} due to: {e}, abort!")
                 raise e
+            finally:
+                if isinstance(data_stream, StreamWrapper):
+                    data_stream.autoclose()
 
     def __len__(self) -> int:
         if self.length == -1:

torchdata/datapipes/iter/util/combining.py

@@ -11,6 +11,8 @@
 from torch.utils.data import functional_datapipe, IterDataPipe, MapDataPipe
 from torch.utils.data.datapipes.utils.common import _check_unpickable_fn
 
+from torchdata.datapipes.utils.janitor import janitor
+
 T_co = TypeVar("T_co", covariant=True)
 
 
@@ -109,6 +111,14 @@ def __iter__(self) -> Iterator:
             else:
                 yield res
 
+        for remaining in ref_it:
+            janitor(remaining)
+
+        # TODO(VItalyFedyunin): This should be Exception or warn when debug mode is enabled
+        if len(self.buffer) > 0:
+            for k, v in self.buffer.items():
+                janitor(v)
+
     def __len__(self) -> int:
         return len(self.source_datapipe)
 

torchdata/datapipes/iter/util/decompressor.py

@@ -97,7 +97,9 @@ def __iter__(self) -> Iterator[Tuple[str, StreamWrapper]]:
         for path, file in self.source_datapipe:
             file_type = self._detect_compression_type(path)
             decompressor = self._DECOMPRESSORS[file_type]
-            yield path, StreamWrapper(decompressor(file))
+            yield path, StreamWrapper(decompressor(file), file, name=path)
+            if isinstance(file, StreamWrapper):
+                file.autoclose()
 
 
 @functional_datapipe("extract")

torchdata/datapipes/iter/util/rararchiveloader.py

@@ -107,7 +107,9 @@ def __iter__(self) -> Iterator[Tuple[str, io.BufferedIOBase]]:
                 inner_path = os.path.join(path, info.filename)
                 file_obj = rar.open(info)
 
-                yield inner_path, StreamWrapper(file_obj)  # type: ignore[misc]
+                yield inner_path, StreamWrapper(file_obj, stream)  # type: ignore[misc]
+            if isinstance(stream, StreamWrapper):
+                stream.autoclose()
 
     def __len__(self) -> int:
         if self.length == -1:

torchdata/datapipes/iter/util/tararchiveloader.py

@@ -67,10 +67,13 @@ def __iter__(self) -> Iterator[Tuple[str, BufferedIOBase]]:
                         warnings.warn(f"failed to extract file {tarinfo.name} from source tarfile {pathname}")
                         raise tarfile.ExtractError
                     inner_pathname = os.path.normpath(os.path.join(pathname, tarinfo.name))
-                    yield inner_pathname, StreamWrapper(extracted_fobj)  # type: ignore[misc]
+                    yield inner_pathname, StreamWrapper(extracted_fobj, data_stream, name=inner_pathname)  # type: ignore[misc]
             except Exception as e:
                 warnings.warn(f"Unable to extract files from corrupted tarfile stream {pathname} due to: {e}, abort!")
                 raise e
+            finally:
+                if isinstance(data_stream, StreamWrapper):
+                    data_stream.autoclose()
 
     def __len__(self) -> int:
         if self.length == -1:

torchdata/datapipes/iter/util/xzfileloader.py

@@ -56,10 +56,13 @@ def __iter__(self) -> Iterator[Tuple[str, BufferedIOBase]]:
             try:
                 extracted_fobj = lzma.open(data_stream, mode="rb")  # type: ignore[call-overload]
                 new_pathname = pathname.rstrip(".xz")
-                yield new_pathname, StreamWrapper(extracted_fobj)  # type: ignore[misc]
+                yield new_pathname, StreamWrapper(extracted_fobj, data_stream)  # type: ignore[misc]
             except Exception as e:
                 warnings.warn(f"Unable to extract files from corrupted xz/lzma stream {pathname} due to: {e}, abort!")
                 raise e
+            finally:
+                if isinstance(data_stream, StreamWrapper):
+                    data_stream.autoclose()
 
     def __len__(self) -> int:
         if self.length == -1:

torchdata/datapipes/iter/util/ziparchiveloader.py

@@ -67,10 +67,13 @@ def __iter__(self) -> Iterator[Tuple[str, BufferedIOBase]]:
                         continue
                     extracted_fobj = zips.open(zipinfo)
                     inner_pathname = os.path.normpath(os.path.join(pathname, zipinfo.filename))
-                    yield inner_pathname, StreamWrapper(extracted_fobj)  # type: ignore[misc]
+                    yield inner_pathname, StreamWrapper(extracted_fobj, data_stream, name=inner_pathname)  # type: ignore[misc]
             except Exception as e:
                 warnings.warn(f"Unable to extract files from corrupted zipfile stream {pathname} due to: {e}, abort!")
                 raise e
+            finally:
+                if isinstance(data_stream, StreamWrapper):
+                    data_stream.autoclose()
             # We are unable to close 'data_stream' here, because it needs to be available to use later
 
     def __len__(self) -> int:

torchdata/datapipes/utils/__init__.py

@@ -7,5 +7,6 @@
 from torch.utils.data.datapipes.utils.common import StreamWrapper
 
 from ._visualization import to_graph
+from .janitor import janitor
 
-__all__ = ["StreamWrapper", "to_graph"]
+__all__ = ["StreamWrapper", "janitor", "to_graph"]

torchdata/datapipes/utils/janitor.py

@@ -0,0 +1,10 @@
+from torchdata.datapipes.utils import StreamWrapper
+
+
+def janitor(obj):
+    """
+    Invokes various `obj` cleanup procedures such as:
+    - Closing streams
+    """
+    # TODO(VitalyFedyunin): We can also release caching locks here to allow filtering
+    StreamWrapper.close_streams(obj)