Repository: Lightning-AI/pytorch-lightning
Pull request: Close profiler when `StopIteration` is raised
## What does this PR do?

<!--
Please include a summary of the change and which issue is fixed.
Please also include relevant motivation and context.
List any dependencies that are required for this change.

If we didn't discuss your PR in Github issues there's a high chance it will not be merged.

The following links the related issue to the PR (https://docs.github.com/en/free-pro-team@latest/github/managing-your-work-on-github/linking-a-pull-request-to-an-issue#linking-a-pull-request-to-an-issue-using-a-keyword)
-->

Fixes #14940 

### Does your PR introduce any breaking changes? If yes, please list them.

<!-- FILL IN or None -->

## Before submitting

- [x] Was this **discussed/approved** via a GitHub issue? (not for typos and docs)
- [x] Did you read the [contributor guideline](https://github.com/Lightning-AI/lightning/blob/master/.github/CONTRIBUTING.md), **Pull Request** section?
- [x] Did you make sure your **PR does only one thing**, instead of bundling different changes together?
- [x] Did you make sure to **update the documentation** with your changes? (if necessary)
- [x] Did you write any **new necessary tests**? (not for typos and docs)
- [x] Did you verify new and **existing tests pass** locally with your changes?
- [x] Did you list all the **breaking changes** introduced by this pull request?
- [x] Did you **update the CHANGELOG**? (not for typos, docs, test updates, or minor internal changes/refactors)

<!-- In the CHANGELOG, separate each item in the unreleased section by a blank line to reduce collisions -->

## PR review

Anyone in the community is welcome to review the PR.
Before you start reviewing, make sure you have read the [review guidelines](https://github.com/Lightning-AI/lightning/wiki/Review-guidelines). In short, see the following bullet-list:

- [x] Is this pull request ready for review? (if not, please submit in draft mode)
- [x] Check that all items from **Before submitting** are resolved
- [x] Make sure the title is self-explanatory and the description concisely explains the PR
- [x] Add labels and milestones (and optionally projects) to the PR so it can be classified

## Did you have fun?

Make sure you had fun coding 🙃


Problem statement:
Profiler doesn't stop action upon epoch end when using IterableDataset with no __len__
### First check

- [X] I'm sure this is a bug.
- [X] I've added a descriptive title to this bug.
- [X] I've provided clear instructions on how to reproduce the bug.
- [X] I've added a code sample.
- [X] I've provided any other important info that is required.

### Bug description

When using an IterableDataset that has no __len__ function implemented, the trainer profiler doesn't remove the action that started when calling next(dataloader) when the first epoch ends and a ValueError is thrown.

### How to reproduce the bug

```python
import torch
from torch.utils.data import Dataset, IterableDataset, DataLoader
from torch.nn import functional as F
from pytorch_lightning import LightningModule, Trainer

class CustomIterableDataset(IterableDataset):
  def __init__(self, data):
        self.data = data


  def __iter__(self):
        return iter(self.data)

class Model(LightningModule):
    def __init__(self):
        super().__init__()
        self.l1 = torch.nn.Linear(1, 1)

    def forward(self, x):
        return torch.relu(self.l1(x.view(x.size(0), -1)))

    def training_step(self, batch, batch_nb):
        x = batch
        y = torch.ones(x.shape)
        loss = F.mse_loss(self(x), y)
        return loss

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=0.02)


data = list(range(12))
data = [torch.tensor([x]).float() for x in data]
dataset = CustomIterableDataset(data)
dataloader = DataLoader(dataset, batch_size=4, num_workers=1)

# Initialize a trainer
trainer = Trainer(
    accelerator="auto",
    devices=1 if torch.cuda.is_available() else None,
    max_epochs=3,
    profiler="simple"
)
model = Model()
trainer.fit(model, dataloader)
```


### Error messages and logs

```

# Error messages and logs here please

```
```
GPU available: False, used: False
TPU available: False, using: 0 TPU cores
IPU available: False, using: 0 IPUs
HPU available: False, using: 0 HPUs

  | Name | Type   | Params
--------------------------------
0 | l1   | Linear | 2     
--------------------------------
2         Trainable params
0         Non-trainable params
2         Total params
0.000     Total estimated model params size (MB)
/home/ziyad/Desktop/ptl_issue_example/venv/lib/python3.9/site-packages/pytorch_lightning/trainer/connectors/data_connector.py:236: PossibleUserWarning: The dataloader, train_dataloader, does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` (try 16 which is the number of cpus on this machine) in the `DataLoader` init to improve performance.
  rank_zero_warn(
Epoch 1: : 0it [00:00, ?it/s, loss=0.741, v_num=5]Traceback (most recent call last):
  File "/home/ziyad/Desktop/ptl_issue_example/example.py", line 47, in <module>
    trainer.fit(model, dataloader)
  File "/home/ziyad/Desktop/ptl_issue_example/venv/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py", line 696, in fit
    self._call_and_handle_interrupt(
  File "/home/ziyad/Desktop/ptl_issue_example/venv/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py", line 650, in _call_and_handle_interrupt
    return trainer_fn(*args, **kwargs)
  File "/home/ziyad/Desktop/ptl_issue_example/venv/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py", line 735, in _fit_impl
    results = self._run(model, ckpt_path=self.ckpt_path)
  File "/home/ziyad/Desktop/ptl_issue_example/venv/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py", line 1166, in _run
    results = self._run_stage()
  File "/home/ziyad/Desktop/ptl_issue_example/venv/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py", line 1252, in _run_stage
    return self._run_train()
  File "/home/ziyad/Desktop/ptl_issue_example/venv/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py", line 1283, in _run_train
    self.fit_loop.run()
  File "/home/ziyad/Desktop/ptl_issue_example/venv/lib/python3.9/site-packages/pytorch_lightning/loops/loop.py", line 200, in run
    self.advance(*args, **kwargs)
  File "/home/ziyad/Desktop/ptl_issue_example/venv/lib/python3.9/site-packages/pytorch_lightning/loops/fit_loop.py", line 271, in advance
    self._outputs = self.epoch_loop.run(self._data_fetcher)
  File "/home/ziyad/Desktop/ptl_issue_example/venv/lib/python3.9/site-packages/pytorch_lightning/loops/loop.py", line 195, in run
    self.on_run_start(*args, **kwargs)
  File "/home/ziyad/Desktop/ptl_issue_example/venv/lib/python3.9/site-packages/pytorch_lightning/loops/epoch/training_epoch_loop.py", line 147, in on_run_start
    _ = iter(data_fetcher)  # creates the iterator inside the fetcher
  File "/home/ziyad/Desktop/ptl_issue_example/venv/lib/python3.9/site-packages/pytorch_lightning/utilities/fetching.py", line 180, in __iter__
    self.prefetching()
  File "/home/ziyad/Desktop/ptl_issue_example/venv/lib/python3.9/site-packages/pytorch_lightning/utilities/fetching.py", line 241, in prefetching
    self._fetch_next_batch(iterator)
  File "/home/ziyad/Desktop/ptl_issue_example/venv/lib/python3.9/site-packages/pytorch_lightning/utilities/fetching.py", line 276, in _fetch_next_batch
    start_output = self.on_fetch_start()
  File "/home/ziyad/Desktop/ptl_issue_example/venv/lib/python3.9/site-packages/pytorch_lightning/utilities/fetching.py", line 229, in on_fetch_start
    self._start_profiler()
  File "/home/ziyad/Desktop/ptl_issue_example/venv/lib/python3.9/site-packages/pytorch_lightning/loops/epoch/training_epoch_loop.py", line 155, in _on_before_fetch
    self.trainer.profiler.start(f"[{self.__class__.__name__}].train_dataloader_next")
  File "/home/ziyad/Desktop/ptl_issue_example/venv/lib/python3.9/site-packages/pytorch_lightning/profilers/simple.py", line 69, in start
    raise ValueError(f"Attempted to start {action_name} which has already started.")
ValueError: Attempted to start [TrainingEpochLoop].train_dataloader_next which has already started.
```

### Important info

```

#- Lightning Component (e.g. Trainer, LightningModule, LightningApp, LightningWork, LightningFlow):
#- PyTorch Lightning Version (e.g., 1.5.0): 1.7.7
#- Lightning App Version (e.g., 0.5.2):
#- PyTorch Version (e.g., 1.10): 1.12.1+cu102
#- Python version (e.g., 3.9): 3.9
#- OS (e.g., Linux): Linux
#- CUDA/cuDNN version: 10.2
#- GPU models and configuration: None
#- How you installed Lightning(`conda`, `pip`, source): pip
#- Running environment of LightningApp (e.g. local, cloud): local
```

### More info

It is pretty clear where the error is coming from. If you go to pytorch_lightning.utilities.fetching under the DataFetcher class you see the following function:
```
def prefetching(self) -> None:
        iterator = self.dataloader_iter
        assert iterator is not None
        for _ in range(self.prefetch_batches):
            try:
                self._fetch_next_batch(iterator)
            except StopIteration:
                # this would only happen when prefetch_batches > the number of batches available and makes
                # fetching_function jump directly to the empty iterator case without trying to fetch again
                self.done = True
                break

```
which calls the _fetch_next_batch function:
```
  def _fetch_next_batch(self, iterator: Iterator) -> None:
      start_output = self.on_fetch_start()
      batch = next(iterator)
      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`
          dataloader = self.dataloader
          assert isinstance(dataloader, Sized)  # `_has_len` is True
          self.done = self.fetched >= len(dataloader)
      self.on_fetch_end(batch, start_output)
```
Upon the end of epoch, inside the _fetch_next_batch function, next(iterator) will raise a StopIteration exception, which will be handled inside the prefetching function, however, self.on_start_fetch starts the profiler and never closes it since self.on_fetch_end is never reached when the exception is raised. Therefore, for the second epoch, self.on_fetch_start() will raise an exception that the profiler already has this action open.

Hints from issue discussion:

Commit message: removed diskcache import

Patch to review:
diff --git a/src/pytorch_lightning/utilities/fetching.py b/src/pytorch_lightning/utilities/fetching.py
index 5dd068af5340d..912ba08f55ec4 100644
--- a/src/pytorch_lightning/utilities/fetching.py
+++ b/src/pytorch_lightning/utilities/fetching.py
@@ -276,14 +276,18 @@ def fetching_function(self) -> Any:
 
     def _fetch_next_batch(self, iterator: Iterator) -> None:
         start_output = self.on_fetch_start()
-        batch = next(iterator)
-        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`
-            dataloader = self.dataloader
-            assert isinstance(dataloader, Sized)  # `_has_len` is True
-            self.done = self.fetched >= len(dataloader)
-        self.on_fetch_end(batch, start_output)
+        try:
+            batch = next(iterator)
+            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`
+                dataloader = self.dataloader
+                assert isinstance(dataloader, Sized)  # `_has_len` is True
+                self.done = self.fetched >= len(dataloader)
+            self.on_fetch_end(batch, start_output)
+        except StopIteration as e:
+            self._stop_profiler()
+            raise e
 
     def move_to_device(self, batch: Any) -> Any:
         if self.store_on_device:


Code review comments:
diff_hunk: @@ -276,14 +276,18 @@ def fetching_function(self) -> Any:
 
     def _fetch_next_batch(self, iterator: Iterator) -> None:
         start_output = self.on_fetch_start()
-        batch = next(iterator)
-        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`
-            dataloader = self.dataloader
-            assert isinstance(dataloader, Sized)  # `_has_len` is True
-            self.done = self.fetched >= len(dataloader)
-        self.on_fetch_end(batch, start_output)
path: src/pytorch_lightning/utilities/fetching.py
text: @user1:
I would try moving the `on_fetch_start` and `on_fetch_end` into `prefetching`, where we already catch the StopIteration

@author:
@user1 Then `on_fetch_start` and `on_fetch_end` would need to be called everywhere where there is `_fetch_next_batch` being called, for example inside `fetching_function` too, you call _fetch_next_batch twice. Here it's only handled once to avoid repetition. What do you think?

@user1:
You are correct. But I would only wrap the `batch = next(iterator)` call, not the entire function.

Can you add a test in `tests/tests_pytorch/utilities/test_fetching.py`? It shouldn't use a trainer etc, just creating this class and faking that a stopiteration is raised with some assertions on a mock profiler would work