Use rich for displaying error traceback in mlinfra cli

pyproject.toml

@@ -34,7 +34,8 @@ dependencies = [
     "GitPython",
     "mypy",
     "getmac",
-    "requests"
+    "requests",
+    "rich"
 ]
 
 [tool.setuptools]

requirements-dev.txt

@@ -34,6 +34,10 @@ jmespath==1.0.1
     # via
     #   boto3
     #   botocore
+markdown-it-py==3.0.0
+    # via rich
+mdurl==0.1.2
+    # via markdown-it-py
 mypy==1.12.0
     # via mlinfra (pyproject.toml)
 mypy-extensions==1.0.0
@@ -48,6 +52,8 @@ pluggy==1.5.0
     # via pytest
 pre-commit==3.7.1
     # via mlinfra (pyproject.toml)
+pygments==2.18.0
+    # via rich
 pytest==8.2.2
     # via
     #   mlinfra (pyproject.toml)
@@ -62,6 +68,8 @@ pyyaml==6.0.2
     #   pre-commit
 requests==2.32.3
     # via mlinfra (pyproject.toml)
+rich==13.9.3
+    # via mlinfra (pyproject.toml)
 s3transfer==0.10.3
     # via boto3
 six==1.16.0

requirements-docs.txt

@@ -50,10 +50,14 @@ markdown==3.6
     #   mkdocs
     #   mkdocs-material
     #   pymdown-extensions
+markdown-it-py==3.0.0
+    # via rich
 markupsafe==2.1.5
     # via
     #   jinja2
     #   mkdocs
+mdurl==0.1.2
+    # via markdown-it-py
 mergedeep==1.3.4
     # via
     #   mkdocs
@@ -89,7 +93,9 @@ pathspec==0.12.1
 platformdirs==4.3.6
     # via mkdocs-get-deps
 pygments==2.18.0
-    # via mkdocs-material
+    # via
+    #   mkdocs-material
+    #   rich
 pymdown-extensions==10.8.1
     # via mkdocs-material
 pyparsing==3.2.0
@@ -118,6 +124,8 @@ requests==2.32.3
     # via
     #   mlinfra (pyproject.toml)
     #   mkdocs-material
+rich==13.9.3
+    # via mlinfra (pyproject.toml)
 s3transfer==0.10.3
     # via boto3
 six==1.16.0

src/mlinfra/cli/__init__.py

@@ -9,3 +9,9 @@
 #     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.
+
+from mlinfra.utils.constants import ENABLE_RICH_TRACEBACK
+from mlinfra.utils.utils import setup_rich_traceback
+
+if ENABLE_RICH_TRACEBACK:
+    setup_rich_traceback()

src/mlinfra/utils/constants.py

@@ -29,3 +29,5 @@
 VERSION: Final = version("mlinfra")
 
 EXECUTABLE_PATH = "/usr/local/bin"
+
+ENABLE_RICH_TRACEBACK = True

src/mlinfra/utils/utils.py

@@ -21,6 +21,9 @@
 import sys
 from typing import List
 
+import invoke
+from rich.traceback import install as rich_tr_install
+
 from .constants import TF_PATH
 
 
@@ -186,3 +189,10 @@ def func_wrapper(*args, **kwargs):  # type: ignore
                 return None
 
     return func_wrapper
+
+
+def setup_rich_traceback(debug_mode: bool = False):
+    if debug_mode:
+        rich_tr_install(show_locals=True)
+    else:
+        rich_tr_install(show_locals=False, suppress=[invoke, subprocess])