[DataLoader2] Ensure finalize and finalize_iteration are called durin…

…g shutdown/exception

ghstack-source-id: eb65f0a99f8dbc6c4b741b0d5d295ed082a08bdc
Pull Request resolved: #846

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