Properly cleanup unclosed files within generator function (#910)

Summary:
There is a case that `file.close` is never called because when generator function has never reached to the end. A simple example would be `zip` two datepipes with different length. The longer DataPipe would never reach the end of generator and then it will be cleaned up by `gc`. So, the line of `file.close` is not executed. (This is the reason that Vitaly has to create this [hack](https://github.com/pytorch/pytorch/blob/4451eb24e6287dff62ff8a7ec0eda6a6998807b0/torch/utils/data/datapipes/iter/combining.py#L573-L583) to retrieve all remaining data to make sure generator function is fully executed)

However, this hack introduces another problem where an infinite datapipe would make `zip` never end as it would try to deplete the infinite iterator. See: #865

So, in this PR, I am adding a `try-finally` clause to make sure the `file.close` is always executed during the destruction of `generator` object. Then, we don't need the hack within `zip` any more.

- pytorch/pytorch#89973
- pytorch/vision#6997
- pytorch/pytorch#89974

Pull Request resolved: #910

Reviewed By: wenleix, NivekT

Differential Revision: D41633909

Pulled By: ejguan

fbshipit-source-id: 5613e131dc8b2962c22d2bc7e3a4bb3039440c48

Author: ejguan

torchdata/datapipes/iter/util/combining.py

@@ -115,13 +115,12 @@ def __iter__(self) -> Iterator:
                 else:
                     yield res
         finally:
-            for remaining in ref_it:
-                janitor(remaining)
-
+            del ref_it
             # TODO(633): This should be Exception or warn when debug mode is enabled
-            if len(self.buffer) > 0:
+            if self.buffer:
                 for _, v in self.buffer.items():
                     janitor(v)
+                self.buffer.clear()
 
     def __len__(self) -> int:
         return len(self.source_datapipe)
@@ -156,7 +155,10 @@ def __setstate__(self, state):
         self.buffer = OrderedDict()
 
     def __del__(self):
-        self.buffer.clear()
+        if self.buffer:
+            for _, v in self.buffer.items():
+                janitor(v)
+            self.buffer.clear()
 
 
 @functional_datapipe("zip_with_map")

torchdata/datapipes/iter/util/decompressor.py

@@ -95,11 +95,13 @@ def _detect_compression_type(self, path: str) -> CompressionType:
 
     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), file, name=path)
-            if isinstance(file, StreamWrapper):
-                file.autoclose()
+            try:
+                file_type = self._detect_compression_type(path)
+                decompressor = self._DECOMPRESSORS[file_type]
+                yield path, StreamWrapper(decompressor(file), file, name=path)
+            finally:
+                if isinstance(file, StreamWrapper):
+                    file.autoclose()
 
 
 @functional_datapipe("extract")

torchdata/datapipes/iter/util/plain_text_reader.py

@@ -41,8 +41,10 @@ def skip_lines(self, file: IO) -> Union[Iterator[bytes], Iterator[str]]:
         with contextlib.suppress(StopIteration):
             for _ in range(self._skip_lines):
                 next(file)
-        yield from file
-        file.close()
+        try:
+            yield from file
+        finally:
+            file.close()
 
     def strip_newline(self, stream: Union[Iterator[bytes], Iterator[str]]) -> Union[Iterator[bytes], Iterator[str]]:
         if not self._strip_newline:
@@ -58,10 +60,9 @@ def strip_newline(self, stream: Union[Iterator[bytes], Iterator[str]]) -> Union[
     def decode(self, stream: Union[Iterator[bytes], Iterator[str]]) -> Union[Iterator[bytes], Iterator[str]]:
         if not self._decode:
             yield from stream
-            return
-
-        for line in stream:
-            yield line.decode(self._encoding, self._errors) if isinstance(line, bytes) else line
+        else:
+            for line in stream:
+                yield line.decode(self._encoding, self._errors) if isinstance(line, bytes) else line
 
     def return_path(self, stream: Iterator[D], *, path: str) -> Iterator[Union[D, Tuple[str, D]]]:
         if not self._return_path:

torchdata/datapipes/iter/util/rararchiveloader.py

@@ -90,26 +90,27 @@ def __iter__(self) -> Iterator[Tuple[str, io.BufferedIOBase]]:
             _PATCHED = True
 
         for data in self.datapipe:
-            validate_pathname_binary_tuple(data)
-            path, stream = data
-            if isinstance(stream, rarfile.RarExtFile) or (
-                isinstance(stream, StreamWrapper) and isinstance(stream.file_obj, rarfile.RarExtFile)
-            ):
-                raise ValueError(
-                    f"Nested RAR archive is not supported by {type(self).__name__}. Please extract outer archive first."
-                )
-
-            rar = rarfile.RarFile(stream)
-            for info in rar.infolist():
-                if info.is_dir():
-                    continue
-
-                inner_path = os.path.join(path, info.filename)
-                file_obj = rar.open(info)
-
-                yield inner_path, StreamWrapper(file_obj, stream, name=path)  # type: ignore[misc]
-            if isinstance(stream, StreamWrapper):
-                stream.autoclose()
+            try:
+                validate_pathname_binary_tuple(data)
+                path, stream = data
+                if isinstance(stream, rarfile.RarExtFile) or (
+                    isinstance(stream, StreamWrapper) and isinstance(stream.file_obj, rarfile.RarExtFile)
+                ):
+                    raise ValueError(
+                        f"Nested RAR archive is not supported by {type(self).__name__}. Please extract outer archive first."
+                    )
+
+                rar = rarfile.RarFile(stream)
+                for info in rar.infolist():
+                    if info.is_dir():
+                        continue
+
+                    inner_path = os.path.join(path, info.filename)
+                    file_obj = rar.open(info)
+                    yield inner_path, StreamWrapper(file_obj, stream, name=path)  # type: ignore[misc]
+            finally:
+                if isinstance(stream, StreamWrapper):
+                    stream.autoclose()
 
     def __len__(self) -> int:
         if self.length == -1: