Fix rich import error in Google Colab (#17156)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Carlos Mocholí <carlossmocholi@gmail.com>

src/lightning/pytorch/CHANGELOG.md

@@ -14,6 +14,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 - Fixed issue where pickling the module instance would fail with a DataLoader error ([#17130](https://github.com/Lightning-AI/lightning/pull/17130))
 - Fixed WandbLogger not showing "best" aliases for model checkpoints when `ModelCheckpoint(save_top_k>0)` is used ([#17121](https://github.com/Lightning-AI/lightning/pull/17121))
+- Fixed the availability check for `rich` that prevented Lightning to be imported in Google Colab ([#17156](https://github.com/Lightning-AI/lightning/pull/17156))
 
 
 ## [2.0.0] - 2023-03-15

src/lightning/pytorch/callbacks/rich_model_summary.py

@@ -14,10 +14,10 @@
 from typing import List, Tuple
 
 from lightning.pytorch.callbacks import ModelSummary
-from lightning.pytorch.utilities.imports import _RICH_AVAILABLE
+from lightning.pytorch.callbacks.progress.rich_progress import _RICH_AVAILABLE
 from lightning.pytorch.utilities.model_summary import get_human_readable_count
 
-if _RICH_AVAILABLE:
+if _RICH_AVAILABLE:  # type: ignore[has-type]
     from rich import get_console
     from rich.table import Table
 

src/lightning/pytorch/utilities/imports.py

@@ -12,11 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """General utilities."""
-import operator
 import sys
 
 import torch
-from lightning_utilities.core.imports import compare_version, package_available, RequirementCache
+from lightning_utilities.core.imports import package_available, RequirementCache
 
 _PYTHON_GREATER_EQUAL_3_8_0 = (sys.version_info.major, sys.version_info.minor) >= (3, 8)
 _PYTHON_GREATER_EQUAL_3_10_0 = (sys.version_info.major, sys.version_info.minor) >= (3, 10)
@@ -27,7 +26,6 @@
 _KINETO_AVAILABLE = torch.profiler.kineto_available()
 _OMEGACONF_AVAILABLE = package_available("omegaconf")
 _POPTORCH_AVAILABLE = package_available("poptorch")
-_RICH_AVAILABLE = package_available("rich") and compare_version("rich", operator.ge, "10.2.2")
 _TORCHVISION_AVAILABLE = RequirementCache("torchvision")
 _LIGHTNING_COLOSSALAI_AVAILABLE = RequirementCache("lightning-colossalai")