Merge pull request #636 from mdekstrand/feature/notebook-logging

Improve logging and progress support

Author: mdekstrand

docs/releases/2025.rst

@@ -50,7 +50,7 @@ New Features (incremental)
 
 *   Many LensKit components (batch running, model training, etc.) now report
     progress the progress API in :mod:`lenskit.logging.progress`, and can be
-    connected to TQDM or Rich.
+    connected to Jupyter or Rich.
 *   Added RBP top-N metric (:pr:`334`).
 *   Added command-line tool to fetch datasets (:pr:`347`).
 

lenskit/lenskit/logging/_console.py

@@ -22,13 +22,12 @@ class ConsoleHandler(Handler):
 
     @property
     def supports_color(self) -> bool:
-        return console.is_terminal and not console.no_color
+        return (console.is_terminal or console.is_jupyter) and not console.no_color
 
     def emit(self, record: LogRecord) -> None:
         try:
             fmt = self.format(record)
-            print(fmt)
-            # console.print(*self._decoder.decode(fmt))
+            console.print(self._decoder.decode_line(fmt))
         except Exception:
             self.handleError(record)
 

lenskit/lenskit/logging/config.py

@@ -14,7 +14,7 @@
 
 import structlog
 
-from ._console import ConsoleHandler, setup_console
+from ._console import ConsoleHandler, console, setup_console
 from .processors import format_timestamp, log_warning, remove_internal
 from .progress import set_progress_impl
 from .tracing import lenskit_filtering_logger
@@ -62,7 +62,7 @@ def notebook_logging(level: int = logging.INFO):
     """
     cfg = LoggingConfig()
     cfg.level = level
-    cfg.set_stream_mode("simple")
+    cfg.progress_backend = "notebook"
     cfg.apply()
 
 
@@ -81,6 +81,7 @@ class LoggingConfig:  # pragma: nocover
 
     level: int = logging.INFO
     stream: Literal["full", "simple", "json"] = "full"
+    progress_backend: Literal["notebook", "rich"] | None = None
     file: Path | None = None
     file_level: int | None = None
     file_format: LogFormat = "json"
@@ -108,6 +109,8 @@ def set_stream_mode(self, mode: Literal["full", "simple", "json"]):
         Configure the standard error stream mode.
         """
         self.stream = mode
+        if mode == "full":
+            self.force_console = True
 
     def set_verbose(self, verbose: bool | int = True):
         """
@@ -157,6 +160,12 @@ def apply(self):
             term = logging.StreamHandler(sys.stderr)
             term.setLevel(self.level)
             proc_fmt = structlog.processors.JSONRenderer()
+        elif console.is_jupyter:
+            term = logging.StreamHandler(sys.stdout)
+            term.setLevel(self.level)
+            proc_fmt = structlog.dev.ConsoleRenderer(
+                colors=self.stream == "full" and not console.no_color
+            )
         else:
             term = ConsoleHandler()
             term.setLevel(self.level)
@@ -207,7 +216,9 @@ def apply(self):
 
         root.setLevel(self.effective_level)
 
-        if self.stream == "full":
+        if self.progress_backend is not None:
+            set_progress_impl(self.progress_backend)
+        elif self.stream == "full":
             set_progress_impl("rich")
 
         warnings.showwarning = log_warning

lenskit/lenskit/logging/progress/_dispatch.py

@@ -1,29 +1,30 @@
-from functools import partial
+import warnings
 from typing import Any, Callable, Literal, overload
 
 from ._base import Progress
 
 _backend: Callable[..., Progress] = Progress
 
 
-@overload
-def set_progress_impl(name: Literal["tqdm"], impl: Callable[..., Any] | None = None, /): ...
 @overload
 def set_progress_impl(name: Literal["rich"]): ...
+@overload
+def set_progress_impl(name: Literal["notebook"]): ...
 def set_progress_impl(name: str | None, *options: Any):
+    """
+    Set the progress bar implementation.
+    """
     global _backend
 
     match name:
-        case "tqdm":
-            from tqdm.autonotebook import tqdm
-
-            from ._tqdm import TQDMProgress
-
-            impl = tqdm
-            if options and options[0]:
-                impl = options[0]
-
-            _backend = partial(TQDMProgress, impl)
+        case "notebook":
+            try:
+                from ._notebook import JupyterProgress
+
+                _backend = JupyterProgress
+            except ImportError:
+                warnings.warn("notebook progress backend needs ipywidgets")
+                _backend = Progress
 
         case "rich":
             from ._rich import RichProgress

lenskit/lenskit/logging/progress/_formats.py

@@ -0,0 +1,5 @@
+def field_format(name: str, fs: str | None):
+    if fs:
+        return "{%s:%s}" % (name, fs)
+    else:
+        return "{%s}" % (name,)

lenskit/lenskit/logging/progress/_notebook.py

@@ -0,0 +1,98 @@
+"""
+Jupyter notebook progress support.
+"""
+
+from __future__ import annotations
+
+from time import perf_counter
+
+import ipywidgets as widgets
+from humanize import metric
+from IPython.display import display
+
+from ._base import Progress
+from ._formats import field_format
+
+__all__ = ["JupyterProgress"]
+
+
+class JupyterProgress(Progress):
+    """
+    Progress logging to Jupyter notebook widgets.
+    """
+
+    widget: widgets.IntProgress
+    text: widgets.Label
+    box: widgets.HBox
+    total: int | None
+    current: int
+    _last_update: float = 0
+    _field_format: str | None = None
+
+    def __init__(
+        self,
+        label: str | None,
+        total: int | None,
+        fields: dict[str, str | None],
+    ):
+        self.current = 0
+        self.total = total
+        if total:
+            self.widget = widgets.IntProgress(value=0, min=0, max=total, step=1)
+        else:
+            self.widget = widgets.IntProgress(value=1, min=0, max=1, step=1)
+            self.widget.bar_style = "info"
+
+        pieces = []
+        if label:
+            pieces.append(widgets.Label(value=label))
+        pieces.append(self.widget)
+
+        self.text = widgets.Label()
+        if total:
+            self.text.value = "0 / {}".format(metric(total))
+        pieces.append(self.text)
+
+        self.box = widgets.HBox(pieces)
+        display(self.box)
+
+        if fields:
+            self._field_format = ", ".join(
+                [f"{name}: {field_format(name, fs)}" for (name, fs) in fields.items()]
+            )
+
+    def update(self, advance: int = 1, **kwargs: float | int | str):
+        """
+        Update the progress bar.
+        """
+        self.current += advance
+        now = perf_counter()
+        if now - self._last_update >= 0.1 or (self.total and self.current >= self.total):
+            self.widget.value = self.current
+            if self.total:
+                if self.total >= 1000:
+                    txt = "{} / {}".format(metric(self.current), metric(self.total))
+                else:
+                    txt = "{} / {}".format(self.current, self.total)
+            else:
+                txt = "{} / ?".format(metric(self.current))
+            if self._field_format:
+                txt += " ({})".format(self._field_format.format(**kwargs))
+            self.text.value = txt
+
+            self._last_update = now
+        # if self._field_format:
+        #     self.tqdm.set_postfix_str(self._field_format.format(kwargs))
+
+    def finish(self):
+        """
+        Finish and clean up this progress bar.  If the progresss bar is used as
+        a context manager, this is automatically called on context exit.
+        """
+        self.box.close()
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *args):
+        self.finish()

lenskit/lenskit/logging/progress/_rich.py

@@ -23,6 +23,7 @@
 from .._console import console, get_live
 from .._proxy import get_logger
 from ._base import Progress
+from ._formats import field_format
 
 _log = get_logger("lenskit.logging.progress")
 _pb_lock = Lock()
@@ -38,24 +39,34 @@ class RichProgress(Progress):
     uuid: UUID
     label: str
     total: int | None
-    fields: dict[str, str | None]
     logger: structlog.stdlib.BoundLogger
+    _field_format: str | None = None
     _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._task = _install_bar(self)
 
+        if fields:
+            self._field_format = ", ".join(
+                [
+                    f"[json.key]{name}[/json.key]: {field_format(name, fs)}"
+                    for (name, fs) in fields.items()
+                ]
+            )
+
     def update(self, advance: int = 1, **kwargs: float | int | str):
+        extra = ""
+        if self._field_format:
+            extra = self._field_format.format(**kwargs)
         if _progress is not None:
-            _progress.update(self._task, advance=advance, **kwargs)  # type: ignore
+            _progress.update(self._task, advance=advance, extra=extra)  # type: ignore
 
     def finish(self):
         _remove_bar(self)
@@ -78,12 +89,13 @@ def _install_bar(bar: RichProgress) -> TaskID | None:
                 RateColumn(),
                 TaskProgressColumn(),
                 TimeRemainingColumn(),
+                TextColumn("{task.fields[extra]}"),
                 console=console,
             )
             live.update(_progress)
 
         _active_bars[bar.uuid] = bar
-        return _progress.add_task(bar.label, total=bar.total)
+        return _progress.add_task(bar.label, total=bar.total, extra="")
 
 
 def _remove_bar(bar: RichProgress):

lenskit/lenskit/logging/progress/_tqdm.py

@@ -229,7 +229,7 @@ def compute(
 
         n = len(outputs)
         _log.info("computing %d listwise metrics for %d output lists", len(lms), n)
-        with item_progress("lists", n) as pb:
+        with item_progress("Measuring", n) as pb:
             for i, (key, out) in enumerate(outputs):
                 list_test = test.lookup_projected(key)
                 if out is None:

lenskit/lenskit/metrics/bulk.py

@@ -65,7 +65,7 @@ def __init__(
             log.debug("persisting function")
             job = worker.WorkerData(func, model)
             job = shm_serialize(job, self.manager)
-            log.info("setting up process pool")
+            log.debug("setting up process pool")
             self.pool = ProcessPoolExecutor(n_jobs, ctx, worker.initalize, (job,))
         except Exception as e:
             self.manager.shutdown()