Close profiler when StopIteration is raised (#14945)

Author: ziyadsheeba

src/pytorch_lightning/CHANGELOG.md

@@ -300,6 +300,10 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - `SaveConfigCallback` instances should only save the config once to allow having the `overwrite=False` safeguard when using `LightningCLI(..., run=False)` ([#14927](https://github.com/Lightning-AI/lightning/pull/14927))
 
 
+
+- Fixed an issue with terminating the trainer profiler when a `StopIteration` exception is raised while using an `IterableDataset` ([#14940](https://github.com/Lightning-AI/lightning/pull/14945))
+
+
 ## [1.7.7] - 2022-09-22
 
 ### Fixed

src/pytorch_lightning/utilities/fetching.py

@@ -276,7 +276,11 @@ def fetching_function(self) -> Any:
 
     def _fetch_next_batch(self, iterator: Iterator) -> None:
         start_output = self.on_fetch_start()
-        batch = next(iterator)
+        try:
+            batch = next(iterator)
+        except StopIteration as e:
+            self._stop_profiler()
+            raise e
         self.fetched += 1
         if not self.prefetch_batches and self._has_len:
             # when we don't prefetch but the dataloader is sized, we use the length for `done`

tests/tests_pytorch/utilities/test_fetching.py

@@ -78,6 +78,31 @@ def generate():
     assert fetcher.fetched == 3
 
 
+@pytest.mark.parametrize("use_combined_loader", [False, True])
+def test_profiler_closing(use_combined_loader):
+    """Tests if the profiler terminates upon raising a StopIteration on an iterable dataset."""
+
+    class TestDataset(IterableDataset):
+        def __init__(self):
+            self.list = list(range(1))
+
+        def __iter__(self):
+            return iter(self.list)
+
+    fetcher = DataFetcher()
+    if use_combined_loader:
+        loader = CombinedLoader([DataLoader(TestDataset()), DataLoader(TestDataset())])
+    else:
+        loader = DataLoader(TestDataset())
+    fetcher.setup(loader)
+    profiler = SimpleProfiler()
+    fetcher._start_profiler = lambda: profiler.start("test")
+    fetcher._stop_profiler = lambda: profiler.stop("test")
+    iter(fetcher)  # on epoch 0 start
+    next(fetcher)  # raises StopIteration exception
+    assert not bool(profiler.current_actions)
+
+
 class EmptyIterDataset(IterableDataset):
     def __iter__(self):
         return iter([])