Repository: Lightning-AI/pytorch-lightning
Pull request: Fix pickling issues with rich progress bar
## What does this PR do?

Fixes #15212

### Does your PR introduce any breaking changes? If yes, please list them.

No.

## 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/PyTorchLightning/pytorch-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 update the [CHANGELOG](https://github.com/PyTorchLightning/pytorch-lightning/blob/master/CHANGELOG.md)? (not for typos, docs, test updates, or internal minor changes/refactorings)

<!-- For CHANGELOG separate each item in the unreleased section by a blank line to reduce collisions -->

## PR review
Anyone in the community is free to review the PR once the tests have passed.
Before you start reviewing make sure you have read [Review guidelines](https://github.com/PyTorchLightning/pytorch-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?
I made sure I had fun coding 🙃


cc @borda @kaushikb11 @rohitgr7

Problem statement:
Unable to pickle (dump) LightningModule object if RichProgressBar callback is added to Trainer
### Bug description

Adding `RichProgressBar` callback to the `Trainer` initialization, and then attempting to save the `LightningModule` object gives a `TypeError`:

```bash
Traceback (most recent call last):
  File "/home/krshrimali/bug.py", line 12, in <module>
    pickle.dumps(boring_model)
TypeError: cannot pickle '_thread.RLock' object
```

After some debugging with the help of @tchaton, here is where I could find the callback object:

```bash
> /home/krshrimali/anaconda3/envs/cloud_weights/lib/python3.9/pickle.py(607)save()
-> self.save_reduce(obj=obj, *rv)
(Pdb) obj
<pytorch_lightning.callbacks.progress.rich_progress.CustomProgress object at 0x7f7c096e7520>
```

(after disabling fast pickling)

If one tries to save the model before `trainer.fit` is called, `pickle.dumps(model)` works. After suggestion with @rohitgr7 - if `model.trainer = None` is done after `trainer.fit`, the error is gone. Not saying that it's a solution, but can help us improve the teardown method.

We should also write tests to make sure it works? `torch.save` is one of the features that everyone will expect to just work with LightningModules and Trainer.

### How to reproduce the bug

```python
from pytorch_lightning.demos.boring_classes import BoringModel
from pytorch_lightning import Trainer
from pytorch_lightning.callbacks import RichProgressBar

import pickle


boring_model = BoringModel()
trainer = Trainer(callbacks=[RichProgressBar()], max_epochs=1)
trainer.fit(model=boring_model)
# Uncomment the following to get rid of the error (not a solution)
# boring_model.trainer = None
pickle.dumps(boring_model)
```


### Error messages and logs

```bash
Traceback (most recent call last):
  File "/home/krshrimali/bug.py", line 12, in <module>
    pickle.dumps(boring_model)
TypeError: cannot pickle '_thread.RLock' object
```

### Environment

```

#- Lightning Component (e.g. Trainer, LightningModule, LightningApp, LightningWork, LightningFlow):
- PyTorch Lightning Version (e.g., 1.5.0): 1.7.1
#- Lightning App Version (e.g., 0.5.2):
#- PyTorch Version (e.g., 1.10):
#- Python version (e.g., 3.9): 3.9.13
#- OS (e.g., Linux): Linux (Pop OS 22.04)
#- CUDA/cuDNN version:
#- GPU models and configuration:
#- How you installed Lightning(`conda`, `pip`, source):
#- Running environment of LightningApp (e.g. local, cloud):

```


### More info

_No response_

Hints from issue discussion:

Commit message: Merge remote-tracking branch 'origin/bugfix/rich-picklable' into bugfix/rich-picklable

Patch to review:
diff --git a/src/pytorch_lightning/CHANGELOG.md b/src/pytorch_lightning/CHANGELOG.md
index 5f687de11a8da..f1d663eaab63d 100644
--- a/src/pytorch_lightning/CHANGELOG.md
+++ b/src/pytorch_lightning/CHANGELOG.md
@@ -12,6 +12,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 ### Changed
 
 
+### Fixed
+
+- Fixed a pickling error when using `RichProgressBar` together with checkpointing ([#15319](https://github.com/Lightning-AI/lightning/pull/15319))
 
 
 ## [1.8.0] - 2022-MM-DD
diff --git a/src/pytorch_lightning/callbacks/progress/rich_progress.py b/src/pytorch_lightning/callbacks/progress/rich_progress.py
index 5677f95a1f9d6..30f616035b59f 100644
--- a/src/pytorch_lightning/callbacks/progress/rich_progress.py
+++ b/src/pytorch_lightning/callbacks/progress/rich_progress.py
@@ -26,7 +26,7 @@
 
 if _RICH_AVAILABLE:
     from rich import get_console, reconfigure
-    from rich.console import RenderableType
+    from rich.console import Console, RenderableType
     from rich.progress import BarColumn, Progress, ProgressColumn, Task, TaskID, TextColumn
     from rich.progress_bar import ProgressBar
     from rich.style import Style
@@ -252,6 +252,7 @@ def __init__(
         super().__init__()
         self._refresh_rate: int = refresh_rate
         self._leave: bool = leave
+        self._console: Optional[Console] = None
         self._console_kwargs = console_kwargs or {}
         self._enabled: bool = True
         self.progress: Optional[Progress] = None
@@ -563,6 +564,13 @@ def configure_columns(self, trainer: "pl.Trainer") -> list:
             ProcessingSpeedColumn(style=self.theme.processing_speed),
         ]
 
+    def __getstate__(self) -> Dict:
+        state = self.__dict__.copy()
+        # both the console and progress object can hold thread lock objects that are not pickleable
+        state["progress"] = None
+        state["_console"] = None
+        return state
+
 
 def _detect_light_colab_theme() -> bool:
     """Detect if it's light theme in Colab."""
diff --git a/tests/tests_pytorch/callbacks/progress/test_rich_progress_bar.py b/tests/tests_pytorch/callbacks/progress/test_rich_progress_bar.py
index 242794213652b..4706589ac6abc 100644
--- a/tests/tests_pytorch/callbacks/progress/test_rich_progress_bar.py
+++ b/tests/tests_pytorch/callbacks/progress/test_rich_progress_bar.py
@@ -11,6 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import pickle
 from collections import defaultdict
 from unittest import mock
 from unittest.mock import DEFAULT, Mock
@@ -416,3 +417,17 @@ def test_rich_progress_bar_padding():
     train_description = progress_bar._get_train_description(current_epoch=0)
     assert "Epoch 0/0" in train_description
     assert len(progress_bar.validation_description) == len(train_description)
+
+
+def test_rich_progress_bar_can_be_pickled():
+    bar = RichProgressBar()
+    trainer = Trainer(fast_dev_run=True, callbacks=[bar], max_steps=1)
+    model = BoringModel()
+
+    pickle.dumps(bar)
+    trainer.fit(model)
+    pickle.dumps(bar)
+    trainer.test(model)
+    pickle.dumps(bar)
+    trainer.predict(model)
+    pickle.dumps(bar)


Code review comments:
diff_hunk: @@ -416,3 +417,17 @@ def test_rich_progress_bar_padding():
     train_description = progress_bar._get_train_description(current_epoch=0)
     assert "Epoch 0/0" in train_description
     assert len(progress_bar.validation_description) == len(train_description)
+
+
+def test_rich_progress_bar_can_be_pickled():
+    bar = RichProgressBar()
+    trainer = Trainer(fast_dev_run=True, callbacks=[bar], max_steps=1)
path: tests/tests_pytorch/callbacks/progress/test_rich_progress_bar.py
text: @user1:
Mind adding the usual boilerplate of `logger=None, enable_checkpointing=False, enable_model_summary=False`...

@author:
I copied the test function from the tqdm tests. I will keep the checkpointing=True since this test is about pickling.

diff_hunk: @@ -416,3 +417,17 @@ def test_rich_progress_bar_padding():
     train_description = progress_bar._get_train_description(current_epoch=0)
     assert "Epoch 0/0" in train_description
     assert len(progress_bar.validation_description) == len(train_description)
+
+
+def test_rich_progress_bar_can_be_pickled():
+    bar = RichProgressBar()
+    trainer = Trainer(fast_dev_run=True, callbacks=[bar], max_steps=1)
path: tests/tests_pytorch/callbacks/progress/test_rich_progress_bar.py
text: @user1:
Doesn't `fast_dev_run` make `max_steps` redundant?

diff_hunk: @@ -416,3 +417,17 @@ def test_rich_progress_bar_padding():
     train_description = progress_bar._get_train_description(current_epoch=0)
     assert "Epoch 0/0" in train_description
     assert len(progress_bar.validation_description) == len(train_description)
+
+
+def test_rich_progress_bar_can_be_pickled():
+    bar = RichProgressBar()
+    trainer = Trainer(fast_dev_run=True, callbacks=[bar], max_steps=1)
+    model = BoringModel()
+
+    pickle.dumps(bar)
+    trainer.fit(model)
+    pickle.dumps(bar)
+    trainer.test(model)
path: tests/tests_pytorch/callbacks/progress/test_rich_progress_bar.py
text: @user1:
How could you forget about validate????!!!

@author:
I copied the test function from the tqdm tests. It was probably written before validate even existed.