From 54229cd33e147b883827158b006e80723a6a03e0 Mon Sep 17 00:00:00 2001
From: Jiarui Fang <fangjiarui123@gmail.com>
Date: Wed, 16 Mar 2022 09:51:15 +0800
Subject: [PATCH] [log] better logging display with rich (#426)

* better logger using rich

* remove deepspeed in zero requirements
---
 colossalai/logging/logging.py      | 3 ++-
 requirements/requirements-test.txt | 2 +-
 requirements/requirements-zero.txt | 1 -
 requirements/requirements.txt      | 1 +
 4 files changed, 4 insertions(+), 3 deletions(-)

diff --git a/colossalai/logging/logging.py b/colossalai/logging/logging.py
index 089308188..1f0ea6c1b 100644
--- a/colossalai/logging/logging.py
+++ b/colossalai/logging/logging.py
@@ -5,11 +5,12 @@ 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'
-logging.basicConfig(level=logging.INFO, format=_FORMAT)
+logging.basicConfig(level=logging.INFO, format=_FORMAT, handlers=[RichHandler()])
 
 
 class DistributedLogger:
diff --git a/requirements/requirements-test.txt b/requirements/requirements-test.txt
index 69b82ff84..3fa37c0f2 100644
--- a/requirements/requirements-test.txt
+++ b/requirements/requirements-test.txt
@@ -1,3 +1,3 @@
 pytest
 rpyc
-matplotlib
\ No newline at end of file
+matplotlib
diff --git a/requirements/requirements-zero.txt b/requirements/requirements-zero.txt
index 816211e72..e69de29bb 100644
--- a/requirements/requirements-zero.txt
+++ b/requirements/requirements-zero.txt
@@ -1 +0,0 @@
-deepspeed
\ No newline at end of file
diff --git a/requirements/requirements.txt b/requirements/requirements.txt
index 469b4d323..a05697b37 100644
--- a/requirements/requirements.txt
+++ b/requirements/requirements.txt
@@ -6,3 +6,4 @@ psutil
 tensorboard
 packaging
 pre-commit
+rich