Improve progress bar display

lenskit/lenskit/logging/progress/_base.py

@@ -1,12 +1,14 @@
 from __future__ import annotations
 
+from typing import Any
+
 
 class Progress:
     """
     Base class for progress reporting.  The default implementations do nothing.
     """
 
-    def __init__(self, *args, **kwargs):
+    def __init__(self, *args: Any, **kwargs: Any):
         pass
 
     def update(self, advance: int = 0, **kwargs: float | int | str):

lenskit/lenskit/logging/progress/_rich.py

@@ -1,18 +1,31 @@
+# pyright: strict
 from __future__ import annotations
 
 from threading import Lock
 from uuid import UUID, uuid4
 
 import structlog
-from rich.console import Group
+from rich.progress import (
+    BarColumn,
+    MofNCompleteColumn,
+    ProgressColumn,
+    SpinnerColumn,
+    Task,
+    TaskID,
+    TaskProgressColumn,
+    TextColumn,
+    TimeRemainingColumn,
+)
 from rich.progress import Progress as ProgressImpl
-from rich.progress import TaskID
+from rich.text import Text
+from typing_extensions import override
 
 from .._console import console, get_live
 from ._base import Progress
 
 _log = structlog.stdlib.get_logger("lenskit.logging.progress")
 _pb_lock = Lock()
+_progress: ProgressImpl | None = None
 _active_bars: dict[UUID, RichProgress] = {}
 
 
@@ -22,59 +35,86 @@
     """
 
     uuid: UUID
+    label: str
     total: int | None
     fields: dict[str, str | None]
     logger: structlog.stdlib.BoundLogger
-    _bar: ProgressImpl
-    _task: TaskID
+    _task: TaskID | None = None
 
     def __init__(self, label: str, total: int | None, fields: dict[str, str | None]):
         super().__init__()
         self.uuid = uuid4()
+        self.label = label
         self.total = total
         self.fields = fields
 
         self.logger = _log.bind(label=label, uuid=str(self.uuid))
 
-        self._bar = ProgressImpl(console=console)
-
-        _install_bar(self)
-
-        self._task = self._bar.add_task(label, total=total)
+        self._task = _install_bar(self)
 
     def update(self, advance: int = 1, **kwargs: float | int | str):
-        self._bar.update(self._task, advance=advance, **kwargs)  # type: ignore
+        if _progress is not None:
+            _progress.update(self._task, advance=advance, **kwargs)  # type: ignore
 
     def finish(self):
-        self._bar.stop()
         _remove_bar(self)
 
 
-def _install_bar(bar: RichProgress):
+def _install_bar(bar: RichProgress) -> TaskID | None:
+    global _progress
     bar.logger.debug("installing progress bar")
     live = get_live()
     if live is None:
-        bar._bar.disable = True
-        return
+        return None
 
     with _pb_lock:
+        if _progress is None:
+            _progress = ProgressImpl(
+                SpinnerColumn(),
+                TextColumn("[progress.description]{task.description}"),
+                BarColumn(),
+                MofNCompleteColumn(),
+                RateColumn(),
+                TaskProgressColumn(),
+                TimeRemainingColumn(),
+                console=console,
+            )
+            live.update(_progress)
+
         _active_bars[bar.uuid] = bar
-        rbs = [b._bar for b in _active_bars.values()]
-        group = Group(*rbs)
-        live.update(group)
+        return _progress.add_task(bar.label, total=bar.total)
 
 
 def _remove_bar(bar: RichProgress):
     live = get_live()
-    if live is None:
+    if live is None or _progress is None:
         return
     if bar.uuid not in _active_bars:
         return
+    if bar._task is None:
+        return
 
     bar.logger.debug("uninstalling progress bar")
 
     with _pb_lock:
+        _progress.remove_task(bar._task)
         del _active_bars[bar.uuid]
 
-        live.update(Group(*[b._bar for b in _active_bars.values()]))
-        live.refresh()
+
+class RateColumn(ProgressColumn):
+    def __init__(self):
+        super().__init__()
+
+    @override
+    def render(self, task: Task):
+        speed = task.finished_speed or task.speed
+        if speed is None:
+            disp = "?"
+        elif speed > 1000:
+            disp = "{:d} it/s".format(int(speed))
+        elif speed > 1:
+            disp = "{:.3g} it/s".format(speed)
+        else:
+            disp = "{:.3g} s/it"
+
+        return Text(disp, "progress.percentage")

lenskit/pyproject.toml

@@ -20,14 +20,14 @@ license = { file = "LICENSE.md" }
 dynamic = ["version"]
 dependencies = [
   "pandas ~=2.0",
-  "structlog >= 23.2",
   "numpy >= 1.24",
   "scipy >= 1.10.0",
   "torch ~=2.1",
-  "rich ~=13.5",
   "threadpoolctl >=3.0",
-  "pydantic ~=2.7",
+  "structlog >= 23.2",
+  "rich ~=13.5",
   "pyzmq >=24",
+  "pydantic ~=2.7",
 ]
 
 [project.optional-dependencies]