Make progress bar prettier (#786)

oulgen · web-flow · commit 2ff64260f5dd · 2025-10-02T16:09:41.000-07:00
diff --git a/helion/autotuner/base_search.py b/helion/autotuner/base_search.py
@@ -29,7 +29,7 @@
 import torch.multiprocessing as mp
 from torch.utils._pytree import tree_flatten
 from torch.utils._pytree import tree_map
-from tqdm.auto import tqdm
+from tqdm.rich import tqdm
 from triton.testing import do_bench
 
 from .. import exc
diff --git a/helion/autotuner/benchmarking.py b/helion/autotuner/benchmarking.py
@@ -4,7 +4,7 @@
 import statistics
 from typing import Callable
 
-from tqdm.auto import tqdm
+from tqdm.rich import tqdm
 from triton import runtime
 
 
diff --git a/pyproject.toml b/pyproject.toml
@@ -21,7 +21,8 @@ dependencies = [
     "typing_extensions>=4.0.0",
     "filecheck",
     "psutil",
-    "tqdm"
+    "tqdm",
+    "rich"
 ]
 
 [project.optional-dependencies]
diff --git a/requirements.txt b/requirements.txt
@@ -5,3 +5,4 @@ filecheck
 expecttest
 numpy
 tqdm
+rich

Original file line number	Diff line number	Diff line change
`@@ -21,7 +21,8 @@ dependencies = [`
`21`	`21`	`"typing_extensions>=4.0.0",`
`22`	`22`	`"filecheck",`
`23`	`23`	`"psutil",`
`24`		`- "tqdm"`
	`24`	`+ "tqdm",`
	`25`	`+ "rich"`
`25`	`26`	`]`
`26`	`27`
`27`	`28`	`[project.optional-dependencies]`
-Original file line number
+Diff line change
 expecttest
 numpy
 tqdm
 +rich