From 8d7271c370b7b45f22922ca1f3048c2a40b7bbd6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Adrian=20W=C3=A4lchli?= Date: Tue, 21 Mar 2023 17:31:50 +0100 Subject: [PATCH] Fix rich import error in Google Colab (#17156) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Carlos Mocholí --- src/lightning/pytorch/CHANGELOG.md | 1 + src/lightning/pytorch/callbacks/rich_model_summary.py | 4 ++-- src/lightning/pytorch/utilities/imports.py | 4 +--- 3 files changed, 4 insertions(+), 5 deletions(-) diff --git a/src/lightning/pytorch/CHANGELOG.md b/src/lightning/pytorch/CHANGELOG.md index 8898f7515ef73..e6c3885380870 100644 --- a/src/lightning/pytorch/CHANGELOG.md +++ b/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 diff --git a/src/lightning/pytorch/callbacks/rich_model_summary.py b/src/lightning/pytorch/callbacks/rich_model_summary.py index 0ec1563d70279..8165822c98236 100644 --- a/src/lightning/pytorch/callbacks/rich_model_summary.py +++ b/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 diff --git a/src/lightning/pytorch/utilities/imports.py b/src/lightning/pytorch/utilities/imports.py index b1b441e9e6827..c8b3ea178e37a 100644 --- a/src/lightning/pytorch/utilities/imports.py +++ b/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")