Cleanup logging (#379)

clean up logging

nerfactory/cuda/backend.py

@@ -16,8 +16,11 @@ import glob
 import os
 from subprocess import DEVNULL, call
 
+from rich import console
 from torch.utils.cpp_extension import load
 
+console = console.Console()
+
 PATH = os.path.dirname(os.path.abspath(__file__))
 
 
@@ -38,12 +41,15 @@ else:
 
 extra_cflags = ["-O2"]
 extra_cuda_cflags = ["-O2"]
-_C = load(
-    name="nerfactory_extension",
-    sources=sources,
-    extra_cflags=extra_cflags,
-    extra_cuda_cflags=extra_cuda_cflags,
-    verbose=True,
-)
+
+with console.status(
+    "[bold yellow]Setting up CUDA (This may take a few minutes the first time)", spinner="bouncingBall"
+):
+    _C = load(
+        name="nerfactory_extension",
+        sources=sources,
+        extra_cflags=extra_cflags,
+        extra_cuda_cflags=extra_cuda_cflags,
+    )
 
 __all__ = ["_C"]

pyproject.toml

@@ -44,7 +44,7 @@ dependencies = [
     "torch==1.12.1",
     "torchmetrics[image]==0.8.2",
     "torchtyping>=0.1.4",
-    "torchvision>=0.12.0",
+    "torchvision>=0.13.0",
     "tqdm==4.64.0",
     "u-msgpack-python>=2.4.1",
     "wandb==0.13.2"

scripts/debugging/run_actions.py

@@ -4,11 +4,15 @@ import sys
 
 import dcargs
 import yaml
+from rich.console import Console
+from rich.style import Style
+
+console = Console(width=120)
 
 LOCAL_TESTS = ["Run license checks", "Run isort", "Run Black", "Python Pylint", "Test with pytest"]
 
 
-def run_command(command: str, continue_on_fail: bool = False) -> None:
+def run_command(command: str, continue_on_fail: bool = False) -> bool:
     """Run a command kill actions if it fails
 
     Args:
@@ -16,9 +20,10 @@ def run_command(command: str, continue_on_fail: bool = False) -> None:
         continue_on_fail: whether to continue running commands if the current one fails. Defaults to False.
     """
     ret_code = subprocess.call(command, shell=True)
-    if not continue_on_fail and ret_code != 0:
-        print(f"\033[31mError: `{command}` failed. Exiting...\033[0m")
-        sys.exit(1)
+    if ret_code != 0:
+        console.print(f"[bold red]Error: `{command}` failed.")
+        if not continue_on_fail:
+            sys.exit(1)
     return ret_code == 0
 
 
@@ -33,8 +38,6 @@ def run_github_actions_file(filename: str, continue_on_fail: bool = False):
         my_dict = yaml.safe_load(f)
     steps = my_dict["jobs"]["build"]["steps"]
 
-    print_green = lambda x: print(f"\033[32m{x}\033[0m")
-    print_red = lambda x: print(f"\033[31m{x}\033[0m")
     success = True
 
     for step in steps:
@@ -43,29 +46,31 @@ def run_github_actions_file(filename: str, continue_on_fail: bool = False):
             compressed = compressed.replace("--check", "")
             curr_command = f"{compressed}"
 
-            print("*" * 100)
-            print_green(f"Running: {curr_command}")
+            console.line()
+            console.rule(f"[bold green]Running: {curr_command}")
             success = success and run_command(curr_command, continue_on_fail=continue_on_fail)
         else:
-            print("*" * 100)
-            print(f"Skipping {step}")
+            skip_name = step["name"] if "name" in step else step["uses"]
+            console.print(f"Skipping {skip_name}")
 
     # Add checks for building documentation
-    print_green("Adding notebook documentation metadata")
-    success = success and run_command("python scripts/docs/add_nb_tags.py")
-    print_green("Building Documentation")
-    success = success and run_command("cd docs/; make html SPHINXOPTS='-W;'")
+    console.line()
+    console.rule("[bold green]Adding notebook documentation metadata")
+    success = success and run_command("python scripts/docs/add_nb_tags.py", continue_on_fail=continue_on_fail)
+    console.line()
+    console.rule("[bold green]Building Documentation")
+    success = success and run_command("cd docs/; make html SPHINXOPTS='-W;'", continue_on_fail=continue_on_fail)
 
     if success:
-        print("\n")
-        print_green("=" * 100)
-        print_green("ALL CHECKS PASSED")
-        print_green("=" * 100)
+        console.line()
+        console.rule(characters="=")
+        console.print("[bold green]:TADA: :TADA: :TADA: ALL CHECKS PASSED :TADA: :TADA: :TADA:", justify="center")
+        console.rule(characters="=")
     else:
-        print("\n")
-        print_red("=" * 100)
-        print_red("ERRORS FOUND")
-        print_red("=" * 100)
+        console.line()
+        console.rule(characters="=", style=Style(color="red"))
+        console.print("[bold red]:skull: :skull: :skull: ERRORS FOUND :skull: :skull: :skull:", justify="center")
+        console.rule(characters="=", style=Style(color="red"))
 
 
 def run_code_checks(continue_on_fail: bool = False):

scripts/docs/add_nb_tags.py

@@ -5,6 +5,9 @@ from glob import glob
 
 import dcargs
 import nbformat as nbf
+from rich.console import Console
+
+console = Console(width=120)
 
 
 def main(check: bool = False):
@@ -28,8 +31,6 @@ def main(check: bool = False):
         "# COLLAPSED": "hide-input",  # Hide the input w/ a button to show
     }
 
-    print_yellow = lambda x: print(f"\033[33m{x}\033[0m")  # print in yellow
-
     # Search through each notebook and look for the text, add a tag if necessary
     any_missing = False
     for ipath in notebooks:
@@ -46,7 +47,7 @@ def main(check: bool = False):
                     found_tags.append(val)
 
             if len(found_keys) > 1:
-                print_yellow(f"Found multiple tags {found_keys} for {ipath}")
+                console.print(f"[bold yellow]Found multiple tags {found_keys} for {ipath}")
                 sys.exit(1)
 
             if len(cell_tags) != len(found_tags):
@@ -58,14 +59,16 @@ def main(check: bool = False):
             cell["metadata"]["tags"] = found_tags
         if incorrect_metadata:
             if check:
-                print_yellow(f"{ipath} has incorrect metadata. Call `python scripts.docs.add_nb_tags.py` to add it.")
+                console.print(
+                    f"[bold yellow]{ipath} has incorrect metadata. Call `python scripts.docs.add_nb_tags.py` to add it."
+                )
                 any_missing = True
             else:
                 print(f"Adding metadata to {ipath}")
                 nbf.write(ntbk, ipath)
 
     if not any_missing:
-        print("All notebooks have correct metadata.")
+        console.print("[green]All notebooks have correct metadata.")
 
     if check and any_missing:
         sys.exit(1)

scripts/docs/build_docs.py

@@ -3,6 +3,10 @@ import subprocess
 import sys
 
 import dcargs
+from rich.console import Console
+from rich.style import Style
+
+console = Console(width=120)
 
 LOCAL_TESTS = ["Run license checks", "Run Black", "Python Pylint", "Test with pytest"]
 
@@ -15,7 +19,7 @@ def run_command(command: str) -> None:
     """
     ret_code = subprocess.call(command, shell=True)
     if ret_code != 0:
-        print(f"\033[31mError: `{command}` failed. Exiting...\033[0m")
+        console.print(f"[bold red]Error: `{command}` failed. Exiting...")
         sys.exit(1)
 
 
@@ -26,22 +30,20 @@ def main(clean_cache: bool = False):
         clean_cache: whether to clean the cache before building docs.
     """
 
-    print_green = lambda x: print(f"\033[32m{x}\033[0m")
-
-    print_green("Adding notebook documentation metadata")
+    console.print("[green]Adding notebook documentation metadata")
     run_command("python scripts/docs/add_nb_tags.py")
 
     # Add checks for building documentation
-    print_green("Building Documentation")
+    console.print("[green]Building Documentation")
     if clean_cache:
         run_command("cd docs/; make clean; make html SPHINXOPTS='-W;'")
     else:
         run_command("cd docs/; make html SPHINXOPTS='-W;'")
 
-    print("\n")
-    print_green("=" * 100)
-    print_green("Done")
-    print_green("=" * 100)
+    console.line()
+    console.rule(characters="=", style=Style(color="green"))
+    console.print("[bold green]Done")
+    console.rule(characters="=", style=Style(color="green"))
 
 
 if __name__ == "__main__":