[install] run with out rich (#513)

81145208 · Jiarui Fang · GitHub · 0f2d2191 · 81145208
Unverified Commit 81145208 authored Mar 24, 2022 by Jiarui Fang Committed by GitHub Mar 24, 2022
Hide whitespace changes
Inline Side-by-side

Showing with 7 additions and 3 deletions

colossalai/logging/logging.py colossalai/logging/logging.py +7 -3

No files found.
--- a/colossalai/logging/logging.py
+++ b/colossalai/logging/logging.py
@@ -5,12 +5,16 @@ import colossalai
 import logging
 from pathlib import Path
 from typing import Union
-from rich.logging import RichHandler
 from colossalai.context.parallel_mode import ParallelMode
-_FORMAT = 'colossalai - %(name)s - %(asctime)s %(levelname)s: %(message)s'
+try:
-logging.basicConfig(level=logging.INFO, format=_FORMAT, handlers=[RichHandler()])
+    from rich.logging import RichHandler
+    _FORMAT = 'colossalai - %(name)s - %(asctime)s %(levelname)s: %(message)s'
+    logging.basicConfig(level=logging.INFO, format=_FORMAT, handlers=[RichHandler()])
+except ImportError:
+    _FORMAT = 'colossalai - %(name)s - %(asctime)s %(levelname)s: %(message)s'
+    logging.basicConfig(level=logging.INFO, format=_FORMAT)
 class DistributedLogger: