Ensure finalize and finalize_iteration are called during shutdown/exc…

…eption (pytorch#846)

Summary:
Pull Request resolved: pytorch#846

Fixes pytorch#821

1. Ensure `finalize_iteration` is called during `DataLoader2.shutdown()`
2. If an uncaught exception is raised by `DataLoader2Iterator.next()`, ensure finalization calls are made

Test Plan: Imported from OSS

Reviewed By: ejguan

Differential Revision: D40569566

Pulled By: NivekT

fbshipit-source-id: f0768a62fcc6dfba1aef3a70c8c55563a45304aa

torchdata/dataloader2/dataloader2.py

@@ -67,6 +67,10 @@ def __next__(self) -> T_co:
                 if self.dataloader.reading_service is not None:
                     self.dataloader.reading_service.finalize_iteration()
                 raise
+            except Exception:
+                if self.dataloader:
+                    self.dataloader.shutdown()
+                raise
         else:
             if self.dataloader.reading_service is not None:
                 self.dataloader.reading_service.finalize_iteration()
@@ -198,6 +202,7 @@ def shutdown(self) -> None:
             self._datapipe_iter = None
         if not self._terminated:
             if self.reading_service is not None:
+                self.reading_service.finalize_iteration()
                 self.reading_service.finalize()
             self._terminated = True